Shifting to JAX (WIP) #23
Conversation
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #23 +/- ##
==========================================
- Coverage 99.05% 98.91% -0.14%
==========================================
Files 20 27 +7
Lines 1689 2305 +616
==========================================
+ Hits 1673 2280 +607
- Misses 16 25 +9 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
I left some comments and questions. A few other items.
- It might be good to add a series of unit tests that ensure that each function in the non-jax code that has a counterpart in the jax code gives the same answer.
- It appears that
dkis compute from image sizes in various places? If yes, then we should be able to compute that as a compile-time constant and remove the need for us to specify it by hand.
|
|
||
|
|
||
| @jax.tree_util.register_pytree_node_class | ||
| class NTObservation(NamedTuple): |
There was a problem hiding this comment.
Is the goal here to make a JAX-compatible copy of the ngmix observation or maybe should we restrict this class to only the items we need for running deep-field metadetect?
There was a problem hiding this comment.
Yes, I agree on that we should remove the unnecessary items. But I was trying ot keep it as close to ngmix as possible until we remove all dependency from ngmix (Gaussian moments, etc...). So maybe for now we can continue with the current class?
| dk_w=2 * jnp.pi / (53 * 0.2) / 4, | ||
| dk_d=2 * jnp.pi / (53 * 0.2) / 4, |
There was a problem hiding this comment.
Can you comment on what the various factors here are? Is the factor of 53 related to the PSF size or the image size or do they happen to be the same?
| dk_w=2 * jnp.pi / (53 * 0.2) / 4, | ||
| dk_d=2 * jnp.pi / (53 * 0.2) / 4, | ||
| nxy=201, | ||
| nxy_psf=53, |
There was a problem hiding this comment.
Again I don't follow these computations.
| @@ -15,13 +14,9 @@ class DFMdetObservation(NamedTuple): | |||
| bmask: Optional[jax.Array] | |||
| ormask: Optional[jax.Array] | |||
| noise: Optional[jax.Array] | |||
| jacobian: Optional[jax.Array] | |||
| aft: Optional[jax_galsim.wcs.AffineTransform] | |||
There was a problem hiding this comment.
Looking ahead, I think we should call this wcs.
Uh oh!
There was an error while loading. Please reload this page.