Time lag computation fails starting with dask v2025.4.0

[wtbarnes]

Describe the bug

Using the (local) distributed scheduler with dask and distributed v2025.4.0 (and above) fails with a UnitConversionError. Presumably this is because the units of seconds are being dropped somewhere in the Dask computation, but time_lag expects the return to be a Quantity with a physical type of time. Interestingly, this does not fail if you don't start a client first (i.e. without using distributed).

This might be related to #260.

To Reproduce

import dask.array
import dask.distributed
import numpy as np
import astropy.units as u
from sunkit_image.time_lag import time_lag

client = dask.distributed.Client()

time = np.linspace(0, 1, 500) * u.s
shape = time.shape + (10,10)
s_a = np.random.rand(*shape)
s_b = np.random.rand(*shape)

s_a = dask.array.from_array(s_a, chunks=s_a.shape[:1] + (5, 5))
s_b = dask.array.from_array(s_b, chunks=s_b.shape[:1] + (5, 5))
tl_map = time_lag(s_a, s_b, time)
_ = tl_map.compute()

raises the following exception

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py:1697](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py#line=1696), in _to_own_unit()
   1696 try:
-> 1697     _value = value.to_value(unit)
   1698 except AttributeError:
   1699     # We're not a Quantity.
   1700     # First remove two special cases (with a fast test):
   1701     # 1) Maybe masked printing? MaskedArray with quantities does not
   1702     # work very well, but no reason to break even repr and str.
   1703     # 2) np.ma.masked? useful if we're a MaskedQuantity.

AttributeError: 'numpy.ndarray' object has no attribute 'to_value'

During handling of the above exception, another exception occurred:

UnitConversionError                       Traceback (most recent call last)
File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py:974](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py#line=973), in to_value()
    973 try:
--> 974     scale = self.unit._to(unit)
    975 except Exception:
    976     # Short-cut failed; try default (maybe equivalencies help).

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/core.py:1168](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/core.py#line=1167), in _to()
   1166         return self_decomposed.scale [/](http://localhost:8888/) other_decomposed.scale
-> 1168 raise UnitConversionError(f"'{self!r}' is not a scaled version of '{other!r}'")

UnitConversionError: 'Unit(dimensionless)' is not a scaled version of 'Unit("s")'

During handling of the above exception, another exception occurred:

UnitConversionError                       Traceback (most recent call last)
Cell In[9], line 13
     10 tl_map = time_lag(s_a, s_b, time)
     11 print(tl_map)
---> 13 foo = tl_map.compute()

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/dask/base.py:373](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/dask/base.py#line=372), in DaskMethodsMixin.compute(self, **kwargs)
    349 def compute(self, **kwargs):
    350     """Compute this dask collection
    351 
    352     This turns a lazy Dask collection into its in-memory equivalent.
   (...)
    371     dask.compute
    372     """
--> 373     (result,) = compute(self, traverse=False, **kwargs)
    374     return result

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/dask/base.py:681](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/dask/base.py#line=680), in compute(traverse, optimize_graph, scheduler, get, *args, **kwargs)
    678     expr = expr.optimize()
    679     keys = list(flatten(expr.__dask_keys__()))
--> 681     results = schedule(expr, keys, **kwargs)
    683 return repack(results)

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py:1896](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py#line=1895), in __array_function__()
   1894 function_helper = FUNCTION_HELPERS[function]
   1895 try:
-> 1896     args, kwargs, unit, out = function_helper(*args, **kwargs)
   1897 except NotImplementedError:
   1898     return self._not_implemented_or_raise(function, types)

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity_helper/function_helpers.py:467](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity_helper/function_helpers.py#line=466), in concatenate()
    463 @function_helper
    464 def concatenate(arrays, axis=0, out=None, **kwargs):
    465     # TODO: make this smarter by creating an appropriately shaped
    466     # empty output array and just filling it.
--> 467     arrays, kwargs, unit, out = _iterable_helper(*arrays, out=out, axis=axis, **kwargs)
    468     return (arrays,), kwargs, unit, out

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity_helper/function_helpers.py:459](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity_helper/function_helpers.py#line=458), in _iterable_helper()
    456 if out is not None:
    457     kwargs["out"] = _quantity_out_as_array(out)  # raises if not Quantity.
--> 459 arrays, unit = _quantities2arrays(*args)
    460 return arrays, kwargs, unit, out

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity_helper/function_helpers.py:447](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity_helper/function_helpers.py#line=446), in _quantities2arrays()
    443 # We use the private _to_own_unit method here instead of just
    444 # converting everything to quantity and then do .to_value(qs0.unit)
    445 # as we want to allow arbitrary unit for 0, inf, and nan.
    446 try:
--> 447     arrays = tuple((q._to_own_unit(arg)) for arg in args)
    448 except TypeError:
    449     raise NotImplementedError

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity_helper/function_helpers.py:447](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity_helper/function_helpers.py#line=446), in <genexpr>()
    443 # We use the private _to_own_unit method here instead of just
    444 # converting everything to quantity and then do .to_value(qs0.unit)
    445 # as we want to allow arbitrary unit for 0, inf, and nan.
    446 try:
--> 447     arrays = tuple((q._to_own_unit(arg)) for arg in args)
    448 except TypeError:
    449     raise NotImplementedError

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py:1713](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py#line=1712), in _to_own_unit()
   1711 try:
   1712     as_quantity = Quantity(value)
-> 1713     _value = as_quantity.to_value(unit)
   1714 except UnitsError:
   1715     # last chance: if this was not something with a unit
   1716     # and is all 0, inf, or nan, we treat it as arbitrary unit.
   1717     if not hasattr(value, "unit") and can_have_arbitrary_unit(
   1718         as_quantity.value
   1719     ):

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py:977](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py#line=976), in to_value()
    974     scale = self.unit._to(unit)
    975 except Exception:
    976     # Short-cut failed; try default (maybe equivalencies help).
--> 977     value = self._to_value(unit, equivalencies)
    978 else:
    979     value = self.view(np.ndarray)

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py:883](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/quantity.py#line=882), in _to_value()
    880     equivalencies = self._equivalencies
    881 if not self.dtype.names or isinstance(self.unit, StructuredUnit):
    882     # Standard path, let unit to do work.
--> 883     return self.unit.to(
    884         unit, self.view(np.ndarray), equivalencies=equivalencies
    885     )
    887 else:
    888     # The .to() method of a simple unit cannot convert a structured
    889     # dtype, so we work around it, by recursing.
    890     # TODO: deprecate this?
    891     # Convert simple to Structured on initialization?
    892     result = np.empty_like(self.view(np.ndarray))

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/core.py:1204](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/core.py#line=1203), in to()
   1202     return UNITY
   1203 else:
-> 1204     return self.get_converter(Unit(other), equivalencies)(value)

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/core.py:1133](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/core.py#line=1132), in get_converter()
   1130             else:
   1131                 return lambda v: b(converter(v))
-> 1133 raise exc

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/core.py:1116](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/core.py#line=1115), in get_converter()
   1114 # if that doesn't work, maybe we can do it with equivalencies?
   1115 try:
-> 1116     return self._apply_equivalencies(
   1117         self, other, self._normalize_equivalencies(equivalencies)
   1118     )
   1119 except UnitsError as exc:
   1120     # Last hope: maybe other knows how to do it?
   1121     # We assume the equivalencies have the unit itself as first item.
   1122     # TODO: maybe better for other to have a `_back_converter` method?
   1123     if hasattr(other, "equivalencies"):

File [~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/core.py:1067](http://localhost:8888/lab/tree/Downloads/~/mambaforge/envs/ar-evolution-survey/lib/python3.12/site-packages/astropy/units/core.py#line=1066), in _apply_equivalencies()
   1064 unit_str = get_err_str(unit)
   1065 other_str = get_err_str(other)
-> 1067 raise UnitConversionError(f"{unit_str} and {other_str} are not convertible")

UnitConversionError: '' (dimensionless) and 's' (time) are not convertible