sclab 0.2.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

sclab/tools/cellflow/pseudotime/timeseries.py

@@ -0,0 +1,226 @@
+from typing import Callable, NamedTuple
+
+import numpy as np
+from numpy.lib.stride_tricks import sliding_window_view
+from numpy.typing import NDArray
+from scipy.signal import find_peaks
+from scipy.sparse import csr_matrix, issparse
+from tqdm.auto import tqdm
+
+from ..utils.interpolate import NDBSpline
+
+
+def periodic_sliding_window(
+    data: NDArray, t: NDArray, window_size: int, fn: Callable[[NDArray], NDArray]
+) -> NDArray:
+    ws = window_size + ((window_size - 1) % 2)
+    window_shape = (ws,) + (1,) * (data.ndim - 1)
+
+    o = np.argsort(t)
+    oo = np.argsort(o)
+
+    d = data[o]
+    dd = [*d[-ws // 2 + 1 :], *d, *d[: ws // 2]]
+
+    windows = sliding_window_view(dd, window_shape=window_shape).squeeze()
+    return fn(windows, axis=-1)[oo]
+
+
+def equalization(
+    times: NDArray,
+    t_range: tuple[float, float],
+    max_bins: int = 200,
+    iterations: int = 1e4,
+    tolerance: float = 0.02,
+) -> NDArray:
+    if not isinstance(times, np.ndarray):
+        raise TypeError("times must be a numpy array")
+
+    if times.ndim != 1:
+        raise ValueError("times must be a 1D array")
+
+    t_min, t_max = t_range
+    t_span = t_max - t_min
+
+    # for sorting the values
+    o = np.argsort(times)
+    # and recovering the original order
+    oo = np.argsort(o)
+
+    alpha = 0.1
+    scale_offset = 1
+
+    rng = np.random.default_rng()
+    scaled_times = times.copy()
+
+    for n_bins in tqdm(np.arange(25, max_bins + 1, 25)):
+        for it in range(int(iterations)):
+            bins = np.linspace(t_min, t_max, n_bins + 1)
+            bins[1:-1] += rng.normal(0, t_span / n_bins / 100, bins[1:-1].size)
+            counts, _ = np.histogram(scaled_times, bins=bins)
+            tmp: NDArray = counts / counts.max()
+            rms = np.sqrt(np.mean((tmp - tmp.mean()) ** 2))
+            if rms < tolerance:
+                break
+
+            scales = counts / counts.max() * alpha + scale_offset
+
+            t = scaled_times[o]
+            tt = []
+            i = 0
+            timepoint = 0.0
+            for start, end, scale in zip(bins[:-1], bins[1:], scales):
+                bin_size = end - start
+                new_size = bin_size * scale
+                while i < t.size and t[i] < end:
+                    new_t = (t[i] - start) * scale + timepoint
+                    tt.append(new_t)
+                    i += 1
+                timepoint += new_size
+
+            tt = np.array(tt)
+            scaled_times = tt[oo] / timepoint * t_span + t_min
+
+        else:
+            cnts_mean, cnts_max, cnts_min = counts.mean(), counts.max(), counts.min()
+            print(
+                f"Failed to converge. RMS: {rms}. "
+                + f"({cnts_mean=:.2f}, {cnts_max=:.2f}, {cnts_min=:.2f})"
+            )
+
+    return scaled_times
+
+
+def fit_trends(
+    X: NDArray | csr_matrix,
+    times: NDArray,
+    t_range: tuple[float, float],
+    periodic: bool,
+    grid_size: int = 128,
+    roughness: float | None = None,
+    zero_weight: float = 0.5,
+    window_width: float | None = None,
+    n_timesteps: int | None = None,
+    timestep_delta: float | None = None,
+    progress: bool = True,
+) -> None:
+    if issparse(X):
+        X = np.ascontiguousarray(X.todense())
+
+    tmin, tmax = t_range
+
+    mask = ~np.isnan(times)
+    t = times[mask]
+    X = X[mask]
+
+    F = NDBSpline(
+        grid_size=grid_size,
+        t_range=t_range,
+        periodic=periodic,
+        zero_weight=zero_weight,
+        roughness=roughness,
+        window_width=window_width,
+    )
+    F.fit(t, X, progress=progress)
+
+    eps = np.finfo(float).eps
+    SNR: NDArray
+    SNR = F(t).var(axis=0) / (X.var(axis=0) + eps)
+    SNR = SNR / SNR.max()
+
+    # x = np.linspace(*t_range, 10001)[:-1]
+    # peak_time = x[np.argmax(F(x), axis=0)]
+
+    if n_timesteps is not None and timestep_delta is not None:
+        raise ValueError("Cannot specify both n_timesteps and timestep_delta")
+    elif n_timesteps is None and timestep_delta is None:
+        # default
+        x = np.linspace(*t_range, 101)
+    elif n_timesteps is not None:
+        x = np.linspace(*t_range, n_timesteps)
+    elif timestep_delta is not None:
+        x = np.arange(tmin, tmax + timestep_delta, timestep_delta)
+
+    Y = F(x)
+
+    return x, Y
+
+
+class SinglePeakResult(NamedTuple):
+    times: NDArray
+    heights: NDArray
+    scores: NDArray
+    info: NDArray
+
+
+def find_single_peaks(
+    X: NDArray,
+    t: NDArray,
+    t_range: tuple[float, float] = (0, 1),
+    grid_size: int = 512,
+    periodic: bool = True,
+    zero_weight: float = 0.2,
+    roughness: float = 2,
+    n_timesteps: int = 201,
+    width_range: tuple[float, float] = (0, 100),
+    score_threshold: float = 2.5,
+    progress: bool = True,
+) -> tuple[NDArray, NDArray]:
+    X = X / np.percentile(X + 1, 99, axis=0, keepdims=True)
+    x, Y = fit_trends(
+        X,
+        t,
+        t_range=t_range,
+        periodic=periodic,
+        grid_size=grid_size,
+        zero_weight=zero_weight,
+        roughness=roughness,
+        n_timesteps=n_timesteps,
+        progress=progress,
+    )
+
+    peak_times = np.full(X.shape[1], np.nan)
+    peak_heights = np.full(X.shape[1], np.nan)
+    peak_scores = np.full(X.shape[1], np.nan)
+    peak_info_data = [{}] * X.shape[1]
+
+    idx_sequence = range(X.shape[1])
+    if progress:
+        idx_sequence = tqdm(idx_sequence)
+
+    for i in idx_sequence:
+        y = Y[:, i]
+        k, info = find_peaks(y, prominence=0.05, width=width_range, height=0)
+        m = np.median(y)
+        s = y[k] / m
+        k = k[s > score_threshold]
+        if len(k) == 1:
+            peak_times[i] = x[k]
+            peak_heights[i] = y[k]
+            peak_scores[i] = np.log2(s[0])
+            peak_info_data[i] = info
+
+    return SinglePeakResult(peak_times, peak_heights, peak_scores, peak_info_data)
+
+
+def piecewise_scaling(
+    times: NDArray,
+    t_range: tuple[float, float],
+    start: float,
+    end: float,
+    new_end: float,
+) -> NDArray:
+    tmin, tmax = t_range
+
+    times_pws = np.full(times.shape, np.nan)
+
+    mask = (times >= tmin) & (times < start)
+    times_pws[mask] = times[mask]
+
+    mask = (times >= start) & (times < end)
+    times_pws[mask] = (times[mask] - start) / (end - start) * (new_end - start) + start
+
+    mask = (times >= end) & (times < tmax)
+    times_pws[mask] = (times[mask] - end) / (tmax - end) * (tmax - new_end) + new_end
+
+    return times_pws

sclab/tools/cellflow/utils/density_nd.py

@@ -0,0 +1,136 @@
+from itertools import product
+from typing import Literal, NamedTuple
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy.integrate import trapezoid
+from scipy.interpolate import BSpline
+from sklearn.neighbors import KernelDensity
+
+from .interpolate import fit_smoothing_spline
+
+
+class DensityResult(NamedTuple):
+    kde: KernelDensity
+    grid_size: int
+    bounds: tuple[tuple[float, float], ...]
+    grid: NDArray
+    density: NDArray
+    scale: float
+    periodic: bool
+
+
+def density_nd(
+    data: NDArray,
+    bandwidth: float | Literal["scott", "silverman"] | None = None,
+    algorithm: Literal["kd_tree", "ball_tree", "auto"] = "auto",
+    kernel: str = "gaussian",
+    metric: str = "euclidean",
+    grid_size: tuple | None = None,
+    max_grid_size: int = 2**5 + 1,
+    periodic: bool = False,
+    bounds: tuple[tuple[float, float], ...] | None = None,
+    normalize: bool = False,
+) -> DensityResult:
+    if data.ndim == 1:
+        data = data.reshape(-1, 1)
+
+    nsamples, ndims = data.shape
+    if bounds is None:
+        assert not periodic, "bounds must be specified if periodic=True"
+        lower, upper = data.min(axis=0), data.max(axis=0)
+        span = upper - lower
+        margins = span / 10
+        bounds = tuple(zip(lower - margins, upper + margins))
+    assert len(bounds) == ndims, "must provide bounds for each dimension"
+
+    if periodic:
+        offsets = np.array(list(product([-1, 0, 1], repeat=ndims)))
+        offsets = offsets * np.diff(bounds).T
+        dat = np.empty((nsamples * 3**ndims, ndims))
+        for i, offset in enumerate(offsets):
+            dat[i * nsamples : (i + 1) * nsamples] = data + offset[None, :]
+    else:
+        dat = data
+
+    if bandwidth is None:
+        bandwidth = np.diff(bounds).max() / 64
+
+    kde = KernelDensity(
+        bandwidth=bandwidth,
+        algorithm=algorithm,
+        kernel=kernel,
+        metric=metric,
+    )
+    kde.fit(dat)
+
+    if grid_size is None:
+        max_span = np.diff(bounds).max()
+        rel_span = np.diff(bounds).flatten() / max_span
+        grid_size = tuple((rel_span * max_grid_size).astype(int))
+
+    grid = np.meshgrid(
+        *[np.linspace(*b, n) for b, n in zip(bounds, grid_size)], indexing="ij"
+    )
+    grid = np.vstack([x.ravel() for x in grid]).T
+    d = np.exp(kde.score_samples(grid))
+
+    if normalize and ndims == 1:
+        scale = trapezoid(d, grid.reshape(-1))
+    elif normalize:
+        # perform simple Riemmann sum for higher dimensions
+        deltas = np.diff(bounds).T / (np.array(grid_size) - 1)
+        tmp = d.reshape(grid_size).copy()
+        for i, s in enumerate(grid_size):
+            # take left corners for the sum
+            tmp = tmp.take(np.arange(s - 1), axis=i)
+        scale = tmp.sum() * np.prod(deltas)
+    else:
+        scale = 1
+
+    d /= scale
+
+    return DensityResult(kde, grid_size, bounds, grid, d, scale, periodic)
+
+
+def fit_density_1d(
+    times: NDArray[np.floating],
+    t_range: tuple[float, float],
+    periodic: bool,
+    bandwidth: float | None = None,
+    algorithm: str = "auto",
+    kernel: str = "gaussian",
+    metric: str = "euclidean",
+    max_grid_size: int = 2**8 + 1,
+    lam: float = 1e-5,
+) -> tuple[DensityResult, BSpline]:
+    tmin, tmax = t_range
+    tspan = tmax - tmin
+
+    times_mask = (tmin <= times) * (times <= tmax)
+    times = times[times_mask]
+
+    if bandwidth is None:
+        bandwidth = tspan / 64
+
+    rslt = density_nd(
+        times.reshape(-1, 1),
+        bandwidth=bandwidth,
+        algorithm=algorithm,
+        kernel=kernel,
+        metric=metric,
+        max_grid_size=max_grid_size,
+        periodic=periodic,
+        bounds=(t_range,),
+        normalize=True,
+    )
+
+    bspl = fit_smoothing_spline(
+        rslt.grid[:, 0],
+        rslt.density,
+        t_range,
+        lam=lam,
+        periodic=periodic,
+    )
+
+    return rslt, bspl

sclab/tools/cellflow/utils/interpolate.py

@@ -0,0 +1,334 @@
+import logging
+from typing import Callable
+
+import numpy as np
+from numpy import asarray, ascontiguousarray, floating, prod
+from numpy import empty as np_empty
+from numpy import float64 as np_float64
+from numpy.typing import NDArray
+from scipy.fft import fft, fftfreq
+from scipy.interpolate import BSpline, _fitpack_impl, make_smoothing_spline
+from tqdm.auto import tqdm
+
+from .smoothen import choose_grid_size, count_data_in_intervals, smoothen_data
+
+try:
+    from scipy.interpolate._dierckx import evaluate_spline
+except ImportError:
+    from scipy.interpolate._bspl import evaluate_spline
+
+
+logger = logging.getLogger(__name__)
+PIX2 = 2 * np.pi
+
+
+def fit_smoothing_spline(
+    x: NDArray[floating],
+    y: NDArray[floating],
+    t_range: tuple[float, float],
+    w: NDArray[floating] | None = None,
+    lam: float | None = None,
+    periodic: bool = False,
+    n_reps: int = 3,
+) -> BSpline:
+    if periodic:
+        assert n_reps % 2 == 1
+
+    o = np.argsort(x)
+    x, y = x[o], y[o]
+    if w is not None:
+        w = w[o]
+
+    tmin, tmax = t_range
+    tspan = tmax - tmin
+
+    if periodic:
+        mask = np.logical_and((x >= tmin), (x < tmax))
+    else:
+        mask = np.logical_and((x >= tmin), (x <= tmax))
+
+    x, y = x[mask], y[mask]
+    if w is not None:
+        w = w[mask]
+    n = x.size
+
+    if periodic:
+        xx = np.concatenate([x + i * tspan for i in range(n_reps)])
+        yy = np.tile(y, n_reps)
+        ww = np.tile(w, n_reps) if w is not None else None
+    else:
+        xx = x
+        yy = y
+        ww = w
+
+    bspl = make_smoothing_spline(xx, yy, ww, lam)
+    t, c, k = bspl.tck
+
+    if periodic:
+        N = n_reps // 2
+        t = t - tspan * N
+        t = t[n * N : -n * N + 1]
+        c = c[n * N : -n * N + 1]
+
+    return BSpline(t, c, k)
+
+
+class NDFourier:
+    def __init__(
+        self,
+        xh: NDArray[floating] | None = None,
+        freq: NDArray[floating] | None = None,
+        t_range: tuple[float, float] | None = None,
+        grid_size: int | None = None,
+        periodic: bool = True,
+        largest_harmonic: int = 5,
+        d: int = 0,
+        zero_weight: float = 1.0,
+        smoothing_fn: Callable = np.average,
+    ) -> None:
+        assert periodic
+        assert t_range is not None
+        assert t_range[0] == 0
+
+        self.tmin, self.tmax = self.t_range = t_range
+        self.tscale = PIX2 / self.tmax
+
+        if xh is not None:
+            assert freq is not None
+            self.n = grid_size + 1
+            self.xh = xh.reshape((xh.shape[0], -1, 1)).copy()
+            self.freq = freq.reshape((freq.shape[0], -1, 1)).copy()
+            self.scaled_freq = 1j * self.freq * self.tscale
+
+        self.grid_size = grid_size
+        self.periodic = periodic
+        self.largest_harmonic = largest_harmonic
+        self.d = d
+        self.zero_weight = zero_weight
+        self.smoothing_fn = smoothing_fn
+
+    def fit(
+        self,
+        t: NDArray[floating],
+        X: NDArray[floating],
+    ) -> "NDFourier":
+        if self.grid_size is None:
+            self.grid_size = choose_grid_size(t, self.t_range)
+
+        t_grid = np.linspace(*self.t_range, self.grid_size + 1)
+        self.X_smooth = smoothen_data(
+            t,
+            X,
+            t_range=self.t_range,
+            t_grid=t_grid,
+            periodic=self.periodic,
+            zero_weight=self.zero_weight,
+            fn=self.smoothing_fn,
+        )
+
+        self.n = n = self.X_smooth.shape[0]
+        self.X_smooth = self.X_smooth.reshape((n, -1))
+
+        xh: NDArray[floating] = fft(self.X_smooth, axis=0)
+        freq: NDArray[floating] = fftfreq(n, d=1 / n)
+
+        mask = np.abs(freq) <= self.largest_harmonic
+        xh = xh[mask]
+        freq = freq[mask]
+
+        self.xh = xh.reshape((xh.shape[0], -1, 1))
+        self.freq = freq.reshape((freq.shape[0], -1, 1))
+        self.scaled_freq = 1j * self.freq * self.tscale
+
+        return self
+
+    def derivative(self, d=1) -> "NDFourier":
+        return NDFourier(
+            self.xh,
+            self.freq,
+            self.t_range,
+            self.grid_size,
+            self.periodic,
+            self.largest_harmonic,
+            d + self.d,
+        )
+
+    def __getitem__(self, key) -> "NDFourier":
+        return NDFourier(
+            self.xh[:, key],
+            self.freq,
+            self.t_range,
+            self.grid_size,
+            self.periodic,
+            self.largest_harmonic,
+            self.d,
+        )
+
+    def __call__(self, x: NDArray[floating], d=0) -> NDArray[floating]:
+        x = asarray(x)
+        x_shape = x.shape
+
+        x = ascontiguousarray(x.ravel(), dtype=np_float64)
+
+        d = d + self.d
+        out: NDArray[floating] = np.real(
+            (self.xh * self.scaled_freq**d * np.exp(self.scaled_freq * x)).sum(axis=0)
+            / self.n
+        )
+        out = out.T
+        out = out.reshape(x_shape + (self.xh.shape[1],))
+
+        return out
+
+
+class NDBSpline:
+    def __init__(
+        self,
+        t: NDArray[floating] | None = None,
+        C: NDArray[floating] | None = None,
+        k: int | None = None,
+        t_range: tuple[float, float] | None = None,
+        grid_size: int | None = None,
+        periodic: bool = False,
+        roughness: float | None = None,
+        zero_weight: float = 1.0,
+        window_width: float | None = None,
+        use_grid: bool = True,
+        weight_grid: bool = False,
+        smoothing_fn: Callable = np.average,
+    ) -> None:
+        if periodic:
+            assert t_range is not None
+            assert t_range[0] == 0
+
+        if t is not None or C is not None or k is not None:
+            assert t is not None
+            assert C is not None
+            assert k is not None
+            self.t = t.copy()
+            self.C = C.reshape((C.shape[0], -1)).copy()
+            self.k = k
+
+        if t_range is not None:
+            self.tmin, self.tmax = self.t_range = t_range
+
+        self.grid_size = grid_size
+        self.periodic = periodic
+        self.window_width = window_width
+        self.use_grid = use_grid
+        self.weight_grid = weight_grid
+        self.roughness = roughness
+        self.zero_weight = zero_weight
+        self.smoothing_fn = smoothing_fn
+
+    def fit(
+        self,
+        t: NDArray[floating],
+        X: NDArray[floating],
+        progress: bool = False,
+    ) -> "NDBSpline":
+        X = X.reshape((X.shape[0], -1))
+        if self.t_range is None:
+            self.tmin, self.tmax = self.t_range = t.min(), t.max()
+
+        if self.grid_size is None:
+            self.grid_size = choose_grid_size(t, self.t_range)
+
+        if self.roughness is None:
+            self.roughness = 1
+
+        if self.use_grid:
+            t_grid: NDArray[floating] = np.linspace(*self.t_range, self.grid_size + 1)
+            self.lam = 1 / self.grid_size / 10**self.roughness
+        else:
+            t_grid = None
+            self.lam = 1 / 10**self.roughness
+        self.X_smooth = smoothen_data(
+            t,
+            X,
+            t_range=self.t_range,
+            t_grid=t_grid,
+            periodic=self.periodic,
+            window_width=self.window_width,
+            zero_weight=self.zero_weight,
+            progress=progress,
+            fn=self.smoothing_fn,
+        )
+
+        if t_grid is not None and self.weight_grid:
+            w = np.zeros(self.X_smooth.shape[0], dtype=float)
+            n = count_data_in_intervals(t, t_grid) + 1
+            if self.periodic:
+                n = np.append(n, n[0])
+            else:
+                n = np.append(n, n[-1])
+            w[n > 1] = 1 / np.log1p(n[n > 1])
+        else:
+            w = None
+
+        iterator = self.X_smooth.T
+        if progress:
+            iterator = tqdm(
+                iterator,
+                bar_format="{desc} {percentage:3.0f}%|{bar}|",
+                desc="Fitting bsplines",
+            )
+
+        fit_t_range = (0, 1)
+        fit_t_grid = np.linspace(0, 1, self.grid_size + 1)
+        fit_t = (t - self.tmin) / (self.tmax - self.tmin)
+        C = []
+        for x in iterator:
+            f = fit_smoothing_spline(
+                fit_t_grid if self.use_grid else fit_t,
+                x,
+                t_range=fit_t_range,
+                w=w,
+                lam=self.lam,
+                periodic=self.periodic,
+            )
+
+            C.append(f.c)
+
+        self.t = f.t.copy()
+        self.t *= self.tmax - self.tmin
+        self.t += self.tmin
+        self.C = np.array(C).T.copy()
+        self.k = 3
+
+        return self
+
+    def derivative(self, d: int = 1) -> "NDBSpline":
+        # pad the c array if needed
+        ct = len(self.t) - len(self.C)
+        if ct > 0:
+            self.C = np.r_[self.C, np.zeros((ct, self.C.shape[1]))]
+        t, C, k = _fitpack_impl.splder((self.t, self.C, self.k), d)
+        return NDBSpline(t, C, k, self.t_range, self.grid_size, self.periodic)
+
+    def __getitem__(self, key) -> "NDBSpline":
+        t = self.t
+        C = self.C[:, key]
+        k = self.k
+        return NDBSpline(t, C, k, self.t_range, self.grid_size, self.periodic)
+
+    def __call__(self, x: NDArray[floating], d: int = 0) -> NDArray[floating]:
+        x = asarray(x)
+        x_shape = x.shape
+
+        x = ascontiguousarray(x.ravel(), dtype=np_float64)
+        if self.periodic:
+            n = self.t.size - self.k - 1
+            x = self.t[self.k] + (x - self.t[self.k]) % (self.t[n] - self.t[self.k])
+
+        out = np_empty((len(x), prod(self.C.shape[1:])), dtype=self.C.dtype)
+
+        if not self.t.flags.c_contiguous:
+            self.t = self.t.copy()
+        if not self.C.flags.c_contiguous:
+            self.C = self.C.copy()
+
+        evaluate_spline(self.t, self.C, self.k, x, d, False, out)
+        out = out.reshape(x_shape + (self.C.shape[1],))
+
+        return out