feat: move slow jumps and jumpfix to so3g #1081

skhrg · 2024-12-26T19:08:18Z

Relies on simonsobs/so3g#196
This is a direct rewrite for the most part, no algorithmic changes.

kwolz

I've looked into the changes. I have only two small comments for my ease of understanding, otherwise it looks good to me.

kwolz · 2025-01-13T08:55:03Z

sotodlib/tod_ops/jumps.py

    if isinstance(jumps, np.ndarray):
        jumps = RangesMatrix.from_mask(np.atleast_2d(jumps))
    elif isinstance(jumps, Ranges):
        jumps = RangesMatrix.from_mask(np.atleast_2d(jumps.mask()))
    if not isinstance(jumps, RangesMatrix):
        raise TypeError("jumps not RangesMatrix or convertable to RangesMatrix")
+    jumps = cast(RangesMatrix, jumps)


This seems redundant to me, but it probably isn't. Why?

Honestly that was just to make my typechecker shut up...

Fair enough!

kwolz · 2025-01-13T09:51:11Z

tests/test_tod_ops.py

@@ -310,6 +311,10 @@ def test_jumpfinder(self):
        heights = heights[heights.nonzero()].ravel()
        self.assertTrue(np.all(np.abs(np.array([10, -13, -8]) - np.round(heights)) < 3))

+        # Check fixing
+        ptp_fix = np.ptp(fixed.ravel()[~jumps.buffer(10).mask().ravel()])
+        self.assertTrue(ptp_fix < 1.1*ptp_orig)


Is this really checking the jump fixing? It seems to me this is just checking that the fixed matrix is not messed up in the places where there were no jumps in the first place...

Yeah so the jump fixing assumes everything within the jump range is the jump so you end up with an offset still inside the jump range (which we ignore since we gapfill those ranges anyways).
The ptp checks if we fixed because after the jump there is a DC offset that will be gone in the fixed TOD.

See below for an example (blue is before fixing, orange after):

If we wanted to not have the errors within the jump ranges in the fixed TOD you just need to lower the buffer size for the ranges matrix (but then your jump could be not in the range). The current range is set by the window size used for jump finding to give a reasonable margin of error.

Right, thanks for the visual explanation, the ptp matter is clear now. Just for my understanding: is jumps.buffer(10).mask() flagging the jump regions plus a buffer region covering the 10 samples before and the 10 samples after the respective jump region?

Yup, upped the buffer in the test to be conservative so that things never fail in CI. But also I set the PRNG key so I think that is unneeded.

mmccrackan

I think this looks fine other than the failing checks.

skhrg · 2025-01-15T16:59:56Z

simonsobs/so3g#196 is now merged so this should be good to go very soon

mhasself · 2025-01-15T17:19:14Z

sotodlib/tod_ops/jumps.py

 from scipy.sparse import csr_array
 from skimage.restoration import denoise_tv_chambolle
 from so3g import (
-    matched_jumps,
-    matched_jumps64,
+    block_minmax,


One way to create some backwards-compatibility is to not import these by name. Just import so3g. Then the code will crash when you try to use those functions, rather than when you try to import them.

As written now, everyone's installation will break, if they try to use anything from sotodlib.tod_ops, unless they've updated so3g.

…ibility

skhrg · 2025-02-11T13:45:42Z

Now that the new so3g is out can someone approve this?

mhasself

Minor comments ... looking forward to seeing this out in the wild.

mhasself · 2025-02-11T14:21:15Z

sotodlib/tod_ops/jumps.py

 from numpy.typing import NDArray
-from pixell.utils import block_expand, block_reduce, moveaxis
+from pixell.utils import moveaxis


Replace the one remaining moveaxis with np.moveaxis.

mhasself · 2025-02-11T14:24:40Z

sotodlib/tod_ops/jumps.py

    x_fixed = x
    if not inplace:
        x_fixed = x.copy()
-    orig_shape = x.shape
-    x_fixed = np.atleast_2d(x_fixed)
+        x_fixed = np.ascontiguousarray(x_fixed)


To reduce unnecessary reallocations, make use of np.asarray. To do these two lines, for example, you can just: x_fixed = np.asarray(x, order='C', copy=True)

mhasself · 2025-02-11T14:32:35Z

sotodlib/tod_ops/jumps.py

+    heights = heights.astype(x.dtype)
+    heights = np.ascontiguousarray(heights)


# This will only copy if needed for dtype/order. heights = np.asarray(heights, dtype=x.dtype, order='C')

mhasself · 2025-02-11T14:42:06Z

sotodlib/tod_ops/jumps.py

-
+    orig_shape = x.shape
+    x = np.atleast_2d(x)
+    x = np.ascontiguousarray(x)


If user requested inplace, and x is not contiguous, then the change won't be in-place. So you need to check for that.

For example

x = np.asarray(x, order='C', copy=(False if inplace else None))

Do we really want to raise a value error if inplace is requested but the array is not contiguous?
I could instead raise a warning, make a copy, and then copy back into the original x later. But that may be babying the user too much.

Actually assuming we want to raise the error in that case is the most correct block something like:

x_fixed = np.asarray(np.atleast_2d(x), order="C", copy=(False if inplace else None))

That way the cases are as follows:

in place and already contiguous -> x_fixed just references x

in place and not contiguous -> ValueError

no in place and already contiguous -> x_fixed is a copy of x

no in place and not contiguous -> x_fixed is a contiguous copy of x

Does that make sense?

W.r.t. "no in place and already contiguous -> x_fixed is a copy of x" -- actually not true; np.asarray(... copy=None) will not copy if it doesn't need to.

Another way to approach this, to minimize copies and simplify logic:

Apply all the conversions you need to get x_fixed, minimizing copies where possible.

Check whether you ended up making a copy, or not -- fixed_is_a_copy = np.may_share_memory(x, x_fixed)

Take different actions depending on (fixed_is_a_copy, inplace).

I'm ok with your suggestion of making a copy and then updating x at the end.

The reason to ValueError would be that "inplace" is sometimes used not just to compress code, but to minimize copying of the data -- it's sort of implied that doing things inplace will use less RAM, less time. But I think printing a warning is sufficient to help super-optimizers tweak things up.

mhasself · 2025-02-12T16:48:19Z

Oh, darn... asarray(copy=...) is numpy>=2 :(

skhrg · 2025-02-12T16:49:54Z

Tests failing with TypeError: asarray() got an unexpected keyword argument 'copy', looks like that was added in numpy 2.0.0 (and I believe the PR for numpy 2.0 support in so3g is not yet ready).

I'll redo some logic to sidestep this... but hopefully in the future we can switch to it

skhrg · 2025-02-12T18:56:57Z

...the problem with devving on both my laptop and site computing is sometimes I don't push all my fixes.
Should pass now.

msilvafe · 2025-02-19T14:38:08Z

@skhrg does this close out Issue #975 ?

skhrg requested review from kwolz and mmccrackan January 8, 2025 23:34

kwolz reviewed Jan 13, 2025

View reviewed changes

mmccrackan reviewed Jan 15, 2025

View reviewed changes

mhasself reviewed Jan 15, 2025

View reviewed changes

skhrg added 4 commits January 30, 2025 11:54

feat: move slow jumps and jumpfix to so3g

6641d4b

feat: add tunable parameter for cleaning up estimated heights

49d7cf5

test: add test of jump fixing

0a39996

fix: import so3g rather than functions within so3g to increase compat…

c6dc69c

…ibility

skhrg force-pushed the more_jump_speedups branch from 7aa611f to c6dc69c Compare January 30, 2025 16:54

skhrg requested review from kwolz, mmccrackan and mhasself February 11, 2025 13:45

mmccrackan approved these changes Feb 11, 2025

View reviewed changes

mhasself requested changes Feb 11, 2025

View reviewed changes

skhrg force-pushed the more_jump_speedups branch 2 times, most recently from 24bd10f to b007892 Compare February 12, 2025 18:34

fix: make less array copies, fix some typing issues, and some cleanup

56dd87d

skhrg force-pushed the more_jump_speedups branch from b007892 to 56dd87d Compare February 12, 2025 18:55

skhrg requested a review from mhasself February 12, 2025 19:07

mhasself approved these changes Feb 12, 2025

View reviewed changes

skhrg merged commit b81e1d3 into master Feb 13, 2025
3 checks passed

skhrg deleted the more_jump_speedups branch February 13, 2025 14:04

msilvafe mentioned this pull request May 5, 2025

Jump Improvements #975

Closed

		heights = heights.astype(x.dtype)
		heights = np.ascontiguousarray(heights)

feat: move slow jumps and jumpfix to so3g #1081

feat: move slow jumps and jumpfix to so3g #1081

Uh oh!

Conversation

skhrg commented Dec 26, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kwolz left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kwolz Jan 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mmccrackan left a comment

Choose a reason for hiding this comment

Uh oh!

skhrg commented Jan 15, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

skhrg commented Feb 11, 2025

Uh oh!

mhasself left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mhasself commented Feb 12, 2025

Uh oh!

skhrg commented Feb 12, 2025

Uh oh!

skhrg commented Feb 12, 2025

Uh oh!

Uh oh!

msilvafe commented Feb 19, 2025

Uh oh!

Uh oh!

skhrg commented Dec 26, 2024 •

edited

Loading

kwolz Jan 13, 2025 •

edited

Loading