Fix: Slow light schedule evaluation times#48
Open
ftavella wants to merge 11 commits into
Open
Conversation
Replace nested np.vectorize approach in __add__, __sub__, and concatenate_at with array-native functions. Each LightSchedule now stores an _array_func that operates on full numpy arrays efficiently. - __init__: detect array-capable functions and store as _array_func; constants use np.full_like instead of scalar lambdas - __add__ / __sub__: combine _array_func references without wrapping in a new np.vectorize call, cutting evaluation from O(n*L) to O(n+L) - concatenate_at: use boolean-index slicing on arrays instead of np.piecewise with pre-computed arrays (which was incorrect for array inputs) - _func kept as np.vectorize for backward compatibility Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Tests that LightSchedule.Chang14 (which uses both concatenate_at and __add__ internally) evaluates 14 days of data at dt=0.005h in under 1 second for all six typical_indoor_lux values from the original use-case. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
When a LightSchedule is passed as a callable to another LightSchedule (e.g. ShiftWork's LightSchedule(total_schedule, period=...)), calling light(0.0) returns a 1-element array. In NumPy 1.25+ with warnings-as-errors, float(1-element-array) raises DeprecationWarning which was caught by the try/except and re-raised as a misleading ValueError rejecting valid input. Fix: use float(np.asarray(test_output).flat[0]) to safely extract a scalar from any numeric type without triggering the deprecation. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Wrap _func scalar extraction so np.vectorize always receives a Python float, even when the underlying callable (e.g. a nested LightSchedule) returns a 1-element array. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
LightSchedule.__call__ always returns a 1-d array; NumPy 2.4 strictly rejects putting that array into a float slot in np.vectorize. All 4 `return schedule(time)` calls in ground_truth_shift_schedule now extract a Python float explicitly. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Member
Author
|
Generating the light schedules of the evening light paper takes ~ 3 min with the old code: and basically nothing with the new code: and the files/schedules we are generating are exactly the same: 
|
Open
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #24
np.vectorizeapproach in__add__,__sub__, andconcatenate_atwith array-native function compositionChanges
Root cause:
__add__wrappedlambda t: fn1(t) + fn2(t)in a newLightSchedule()call, which applied anothernp.vectorizeon top. Withnpulses this creatednlevels of vectorize nesting — each level iterating element-by-element in Python — giving O(n×L) scalar function calls for a time array of length L.Fix:
LightSchedule.__init__now stores_array_func: an array-native function that handles full numpy arrays efficiently. Forfrom_pulseschedules this is thenp.piecewise-based function directly; for constants it usesnp.full_like; for arbitrary callables it detects array capability and falls back tonp.vectorizeonly when needed.__add__and__sub__combine_array_funcreferences with a simple lambda (fn1(t) + fn2(t)) and bypassLightSchedule.__init__, avoiding any re-wrapping innp.vectorize. Evaluation is now O(n) array-level calls.concatenate_atuses boolean-index slicing (result[mask] = fn(t[mask])) instead ofnp.piecewisewith pre-evaluated arrays, making it array-native and correct._func(thenp.vectorize-wrapped version) is kept on all objects for backward compatibility.Tests
Added a new test in
nbs/test/test_lights.ipynb:Checklist