adding jax code #30
adding jax code #30
Conversation
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #30 +/- ##
==========================================
- Coverage 99.05% 96.94% -2.11%
==========================================
Files 20 28 +8
Lines 1689 2421 +732
==========================================
+ Hits 1673 2347 +674
- Misses 16 74 +58 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
| - ngmix | ||
| - numba | ||
| - numpy | ||
| - jax<0.7.0 |
There was a problem hiding this comment.
Can you comment on this restriction?
beckermr
left a comment
beckermr
left a comment
There was a problem hiding this comment.
Some minor changes, but looks great!
| return True | ||
|
|
||
|
|
||
| def ngmix_obs_to_dfmd_obs(obs: ngmix.observation.Observation) -> DFMdetObservation: |
| ) | ||
|
|
||
|
|
||
| def dfmd_obs_to_ngmix_obs(dfmd_obs) -> Observation: |
|
|
||
|
|
||
| def get_gauss_reconv_psf_galsim(psf, step=DEFAULT_STEP, flux=1): | ||
| def get_gauss_reconv_psf_galsim(psf, step=DEFAULT_STEP, flux=1, dk=None, kim_size=None): |
|
|
||
| def _render_psf_and_build_obs(image, obs, reconv_psf, weight_fac=1): | ||
| def _render_psf_and_build_obs( | ||
| image, obs, reconv_psf, weight_fac=1, max_min_fft_size=None |
There was a problem hiding this comment.
Let's call this max_min_fft_size just fft_size since it is used to set the FFT size directly. We should make this change everywhere.
| if dfmd_obs1.has_bmask() or dfmd_obs2.has_bmask(): | ||
| new_bmask = _extract_attr(dfmd_obs1, "bmask", jnp.int32) | _extract_attr( | ||
| dfmd_obs2, "bmask", jnp.int32 | ||
| ) | ||
|
|
||
| if dfmd_obs1.has_ormask() or dfmd_obs2.has_ormask(): | ||
| new_ormask = _extract_attr(dfmd_obs1, "ormask", jnp.int32) | _extract_attr( | ||
| dfmd_obs2, "ormask", jnp.int32 | ||
| ) | ||
|
|
||
| if dfmd_obs1.has_noise() or dfmd_obs2.has_noise(): | ||
| new_noise = _extract_attr(dfmd_obs1, "noise") + _extract_attr( | ||
| dfmd_obs2, "noise" | ||
| ) |
There was a problem hiding this comment.
How does JIT properly handle these if statements. Have you tested that?
| ) | ||
|
|
||
| # now add in noise corr to make it match the wide noise | ||
| # TODO: is it after to vextorize? |
There was a problem hiding this comment.
Do you want to address this item?
| force_maxk_psf=0.0, | ||
| max_min_fft_size=1024, | ||
| ): | ||
| """Do metacalibration for a combination of wide+deep datasets.""" |
There was a problem hiding this comment.
This function should have a complete doc string.
| mcal_res[k] = dfmd_obs_to_ngmix_obs(mcal_res[k]) | ||
| mcal_res[k].psf.galsim_obj = reconv_psf | ||
|
|
||
| return mcal_res, kinfo |
There was a problem hiding this comment.
We should return only one value if we can.
| force_maxk_psf=0.0, | ||
| max_min_fft_size=1024, | ||
| ) -> dict: | ||
| """Run deep-field metadetection for a simple scenario of a single band |
Co-authored-by: Matthew R. Becker <[email protected]>
2 participants
No description provided.