diff --git a/README.md b/README.md index 6163925..8cd6fb0 100755 --- a/README.md +++ b/README.md @@ -73,7 +73,9 @@ If you need GPU, use the platform-specific instructions in [GPU Acceleration](#g ## Quickstart -SIMPL follows sklearn conventions: configure hyperparameters at init, pass data to `fit()`. For hyperparameter units, see the Model / Maths section. +SIMPL follows sklearn conventions: configure hyperparameters at init, pass data to `fit()`. +Init parameters (`kernel_bandwidth` etc.) can also be left as `"auto"` and SIMPL will infer sensible values from the data. Typically this is not recommended. See +[Automatic parameter inference](#automatic-parameter-inference). ```python from simpl import SIMPL @@ -213,6 +215,24 @@ $K$ is a Gaussian kernel with bandwidth `kernel_bandwidth`. The denominator corr ### Units and Discretisation All hyperparameters (e.g. `speed_prior`, `kernel_bandwidth`, `bin_size`) are defined in _data units_ (e.g. typically [m/s], [m], [m] but these depend on _your_ data of course), not arbitrary time/spatial-bin units. + +#### Automatic parameter inference + +Scale-dependent parameters default to `"auto"` and are inferred from `Xb` during `fit()`: + +- `bin_size = "auto"`: 1/25 of the largest environment span. +- `kernel_bandwidth = "auto"`: the multivariate Scott bandwidth, no smaller than `bin_size`. +- `speed_prior = "auto"`: the mean behavioral speed. + +A maximally minimal usage now becomes: +```python +model = SIMPL() +model.fit(Y, Xb, time) +``` + +> **Warning:** If neural dynamics are faster than measured behavior, set a larger explicit +> `speed_prior`. +> In general we recommend putting thought into these parameter choices and not just blindly using `"auto"`. diff --git a/examples/simpl_demo.ipynb b/examples/simpl_demo.ipynb index f2c6287..10036ce 100755 --- a/examples/simpl_demo.ipynb +++ b/examples/simpl_demo.ipynb @@ -26,7 +26,7 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -49,8 +49,11 @@ "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", + "import simpl\n", "from simpl import SIMPL, load_demo_data, train_test_split\n", - "from simpl.utils import coarsen_dt, find_time_jumps" + "from simpl.utils import coarsen_dt, find_time_jumps\n", + "\n", + "print(f\"SIMPL version: {simpl.__version__}\")" ] }, { diff --git a/src/simpl/simpl.py b/src/simpl/simpl.py index 13406e3..6d588e8 100755 --- a/src/simpl/simpl.py +++ b/src/simpl/simpl.py @@ -34,8 +34,8 @@ class SIMPL: def __init__( self, # Model hyperparameters - kernel_bandwidth: float = 0.04, - speed_prior: float | None = 1.0, + kernel_bandwidth: float | Literal["auto"] = "auto", + speed_prior: float | Literal["auto"] | None = "auto", behavior_prior: float | None = None, # Environment parameters is_1D_angular: bool = False, @@ -85,17 +85,25 @@ def __init__( Parameters ---------- - kernel_bandwidth : float, optional + kernel_bandwidth : float or "auto", optional The bandwidth of the Gaussian kernel (in the same units as the latent space, e.g. meters) used for KDE when fitting receptive fields. Smaller values give sharper - fields but are noisier; larger values smooth more. By default 0.04. - speed_prior : float or None, optional + fields but are noisier; larger values smooth more. ``"auto"`` defaults to the + multivariate Scott bandwidth estimated from ``Xb``, bounded below by the resolved + ``bin_size``. By default ``"auto"``. + speed_prior : float, "auto", or None, optional Prior on agent speed in units of meters per second. This controls the strength of the Kalman smoother: a low speed prior constrains the decoded trajectory to be smooth, while a high value lets the trajectory follow the spike likelihood more - closely. Set to None to disable Kalman smoothing and let the trajectory follow - the per-bin maximum-likelihood estimate independently in each time bin. By default - 1.0 m/s. + closely. ``"auto"`` uses the mean speed of the behavioral trajectory. + + .. warning:: + For many neural datasets with high-speed dynamics, this automatically inferred + value will be too slow. Set an explicit, larger ``speed_prior`` when the latent + dynamics are expected to evolve faster than measured behavior. + + Set to None to disable Kalman smoothing and let the trajectory follow the per-bin + maximum-likelihood estimate independently in each time bin. By default ``"auto"``. behavior_prior : float or None, optional Prior on how far the latent positions can deviate from the behavioral positions, in units of meters. This acts as a soft constraint pulling the decoded trajectory @@ -912,7 +920,7 @@ def _E_step(self, Y: jax.Array, F: jax.Array) -> dict: ) # Manifold alignment (fit-time only) - self._substatus("E··· aligning") + self._substatus("decode··· aligning") align_dict = {} if self.align_mode_ == "fields": current_peaks = utils.get_field_peaks(F, self.xF_) @@ -959,12 +967,12 @@ def kde_func(mask): trajectory=X, spikes=Y, kernel=kde.gaussian_kernel, - kernel_bandwidth=self.kernel_bandwidth, + kernel_bandwidth=self.kernel_bandwidth_, mask=mask, return_position_density=True, ) - self._substatus("E✓·M tuning curves") + self._substatus("decode✓ · fit··· tuning curves") F, PX = kde_func(self.spike_mask_) FX = self._interpolate_firing_rates(X, F) return {"F": F, "FX": FX, "PX": PX} @@ -1013,7 +1021,7 @@ def _decode( store_log_maps = getattr(self, "save_full_history_", False) # Likelihood maps and Gaussian observation fits (batched internally) - self._substatus("E··· likelihood") + self._substatus("decode··· likelihood") obs = kde.decode_observations( self.xF_, Y, @@ -1042,7 +1050,7 @@ def _decode( ) # Single-pass filter and smooth - self._substatus("E··· kalman filter") + self._substatus("decode··· kalman filter") mu_f, sigma_f = self.kalman_filter_.filter( mu0=mu0_all[0], sigma0=sigma0_all[0], @@ -1053,7 +1061,7 @@ def _decode( mu0_all=mu0_all, sigma0_all=sigma0_all, ) - self._substatus("E··· kalman smooth") + self._substatus("decode··· kalman smooth") mu_s, sigma_s = self.kalman_filter_.smooth( mus_f=mu_f, sigmas_f=sigma_f, @@ -1266,7 +1274,7 @@ def _init_from_data(self, Y, Xb, time, trial_boundaries, align_to_behavior) -> N if self.is_temporal_: time = np.asarray(time, dtype=float) else: - if self.speed_prior is not None: + if self.speed_prior not in (None, "auto"): warnings.warn( "time=None was passed, so SIMPL is treating the data as non-temporal. " "Kalman smoothing is disabled and speed_prior is ignored. " @@ -1335,6 +1343,18 @@ def _init_from_data(self, Y, Xb, time, trial_boundaries, align_to_behavior) -> N ) self.bin_size_ = self.environment_.bin_size + if isinstance(self.kernel_bandwidth, str): + if self.kernel_bandwidth != "auto": + raise ValueError("kernel_bandwidth must be 'auto' or a positive finite number") + self.kernel_bandwidth_ = max(utils._estimate_kernel_bandwidth(Xb), self.bin_size_) + else: + self.kernel_bandwidth_ = self.kernel_bandwidth + if ( + not np.isscalar(self.kernel_bandwidth_) + or not np.isfinite(self.kernel_bandwidth_) + or self.kernel_bandwidth_ <= 0 + ): + raise ValueError("kernel_bandwidth must be 'auto' or a positive finite number") if self.D_ != self.environment_.D: raise ValueError(f"Data has {self.D_} dimensions but environment has {self.environment_.D}") @@ -1367,21 +1387,6 @@ def _init_from_data(self, Y, Xb, time, trial_boundaries, align_to_behavior) -> N self.xF_shape_ = self.environment_.discrete_env_shape self.N_bins_ = len(self.xF_) - # ── Check speed prior against data ── - displacements = np.sqrt(np.sum(np.diff(Xb, axis=0) ** 2, axis=1)) - median_speed = float(np.median(displacements / self.dt_)) - if ( - self.is_temporal_ - and self.speed_prior is not None - and median_speed > 0 - and self.speed_prior < 0.2 * median_speed - ): - warnings.warn( - f"speed_prior ({self.speed_prior:.4g}) is much slower than the median behavioural speed " - f"({median_speed:.4g}). This may impede the decoded trajectory. " - f"Consider increasing speed_prior (e.g. to {median_speed:.2g} or higher)." - ) - # ── Set up Kalman filter, masks, alignment, coordinates ── self._init_infrastructure(trial_boundaries, align_to_behavior) @@ -1433,17 +1438,36 @@ def _init_infrastructure( self.trial_boundaries_, self.trial_slices_, _, _ = self._validate_trial_boundaries( trial_boundaries, self.T_, device ) + self.block_size_ = max(1, int(np.ceil(self.speckle_block_size_seconds / self.dt_))) + if spike_mask is None and self.block_size_ >= self.T_: + raise ValueError( + "speckle_block_size_seconds must be shorter than the recording duration so both train and " + f"validation observations remain available (got block_size={self.block_size_} bins for T={self.T_})" + ) + + if isinstance(self.speed_prior, str) and self.speed_prior != "auto": + raise ValueError("speed_prior must be 'auto', None, or a positive finite number") + if not self.is_temporal_: + self.speed_prior_ = None + elif self.speed_prior == "auto": + behavior = np.asarray(jax.device_get(self.Xb_)) + displacement = np.diff(behavior, axis=0) + if self.is_1D_angular: + displacement = (displacement + np.pi) % (2 * np.pi) - np.pi + self.speed_prior_ = float( + np.mean(np.linalg.norm(displacement, axis=1) / np.diff(np.asarray(jax.device_get(self.time_)))) + ) + else: + self.speed_prior_ = self.speed_prior + if self.speed_prior_ is not None and ( + not np.isscalar(self.speed_prior_) or not np.isfinite(self.speed_prior_) or self.speed_prior_ <= 0 + ): + raise ValueError("speed_prior must be 'auto', None, or a positive finite number") self._init_kalman_filter() - self.block_size_ = max(1, int(np.ceil(self.speckle_block_size_seconds / self.dt_))) if spike_mask is not None: self.spike_mask_ = jax.device_put(np.asarray(spike_mask, dtype=bool), device) else: - if self.block_size_ >= self.T_: - raise ValueError( - "speckle_block_size_seconds must be shorter than the recording duration so both train and " - f"validation observations remain available (got block_size={self.block_size_} bins for T={self.T_})" - ) self.spike_mask_ = utils.create_speckled_mask( size=(self.T_, self.N_neurons_), sparsity=self.val_frac, @@ -1489,7 +1513,7 @@ def _init_kalman_filter(self) -> None: self.speed_prior_requested_ = self.speed_prior self.kalman_off_speed_prior_ = 1e10 speed_prior_effective = ( - self.speed_prior if self.is_temporal_ and self.speed_prior is not None else self.kalman_off_speed_prior_ + self.speed_prior_ if self.is_temporal_ and self.speed_prior_ is not None else self.kalman_off_speed_prior_ ) self.speed_prior_effective_ = speed_prior_effective speed_sigma = speed_prior_effective * self.dt_ @@ -1521,7 +1545,12 @@ def _init_kalman_filter(self) -> None: # Display # ────────────────────────────────────────────────────────────────────────── - _TABLE_HEADER = f" {'iteration':>9} {'status':<20} {'bits-per-spike (train / val)':>36}" + _STATUS_WIDTH = 30 + _METRIC_WIDTH = 20 + _TABLE_HEADER = ( + f" {'iteration':>9} {'status':<{_STATUS_WIDTH}} " + f"{'train bits per spike':>{_METRIC_WIDTH}} {'val bits per spike':>{_METRIC_WIDTH}} " + ) _TABLE_WIDTH = len(_TABLE_HEADER) @staticmethod @@ -1547,16 +1576,18 @@ def _print_row(self, suffix: str = "") -> None: bps_val = float(self.loglikelihoods_.bits_per_spike_val.sel(iteration=e).values) arrow = " " - status = " M✓" if e == 0 else "E✓·M✓" + status = "fit✓" if e == 0 else "decode✓ · fit✓" if e > 0: prev_bps_val = float(self.loglikelihoods_.bits_per_spike_val.sel(iteration=e - 1).values) arrow = " ↑" if bps_val > prev_bps_val else " ↓" val_ll = float(self.loglikelihoods_.logPYXF_val.sel(iteration=e).values) if val_ll < float(self.loglikelihoods_.logPYXF_val.sel(iteration=0).values): - status = "E✓·M✓ !bps9} {status + suffix:<20} {bps_str:>29}" + row = ( + f" {e:>9} {status + suffix:<{self._STATUS_WIDTH}} " + f"{bps_train:>{self._METRIC_WIDTH}.3f} {bps_val:>{self._METRIC_WIDTH}.3f}{arrow}" + ) line = f"\r{row:<{self._TABLE_WIDTH}}" print(line[: self._term_width() + 1], flush=True) # +1 for \r @@ -1575,7 +1606,6 @@ def _print_header(self) -> None: ) mean_fr = total_spikes / duration / self.N_neurons_ empty_frac = float(jnp.mean(jnp.sum(self.Y_, axis=1) == 0)) * 100 - n_trials = len(self.trial_boundaries_) line1 = [ f"{self.N_neurons_} neurons", f"{spike_str} spikes", @@ -1583,15 +1613,26 @@ def _print_header(self) -> None: f"empty time-bins={empty_frac:.0f}%", ] line2 = [ - f"{self.D_}D", f"env-grid ({grid_str})", f"{duration:.1f}s (dt={self.dt_:.2g}s)", - f"n_trials={n_trials}", + ] + + def _parameter(name, requested, effective): + value = "None" if effective is None else f"{effective:.3f}" + suffix = " (auto)" if requested == "auto" else "" + return f"{name}={value}{suffix}" + + line2.append(_parameter("bin_size", self.bin_size, self.bin_size_)) + line3 = [ + _parameter("kernel_bandwidth", self.kernel_bandwidth, self.kernel_bandwidth_), + _parameter("speed_prior", self.speed_prior, self.speed_prior_), + _parameter("behavior_prior", self.behavior_prior, self.behavior_prior), ] title = f"━━ SIMPL ━━━━━ {self._device_str} " print(f"{title}{'━' * (self._TABLE_WIDTH - len(title))}") print(" · ".join(line1)) - print(" · ".join(line2), end="", flush=True) + print(" · ".join(line2)) + print(" · ".join(line3), end="", flush=True) def _print_summary(self) -> None: """Print the end-of-fitting summary with percentage changes.""" @@ -1894,8 +1935,8 @@ def _build_dataset_attrs(self, trial_boundaries) -> dict: "dt": self.dt_, "is_temporal": int(self.is_temporal_), "trial_boundaries": trial_boundaries, - "kernel_bandwidth": self.kernel_bandwidth, - "speed_prior": np.nan if self.speed_prior is None or not self.is_temporal_ else self.speed_prior, + "kernel_bandwidth": self.kernel_bandwidth_, + "speed_prior": np.nan if self.speed_prior_ is None else self.speed_prior_, "behavior_prior": np.nan if self.behavior_prior is None else self.behavior_prior, "is_1D_angular": int(self.is_1D_angular), "align_mode": self.align_mode_ or "none", diff --git a/src/simpl/utils.py b/src/simpl/utils.py index fd98366..e758653 100644 --- a/src/simpl/utils.py +++ b/src/simpl/utils.py @@ -362,6 +362,26 @@ def _circular_conv_fft_1d(x: jax.Array, k: jax.Array) -> jax.Array: # ────────────────────────────────────────────────────────────────────────────── +def _estimate_kernel_bandwidth(X: np.ndarray) -> float: + r"""Estimate an isotropic bandwidth using the multivariate Scott rule. + + $$ + h = \left(\frac{4}{D + 2}\right)^{\frac{1}{D + 4}} + n^{-\frac{1}{D + 4}} + \left(\prod_{d=1}^{D}\sigma_d\right)^{\frac{1}{D}}, + $$ + + where $n$ is the number of samples, $D$ is the latent dimensionality, + and $\sigma_d$ is the sample standard deviation along dimension $d$. + """ + scales = np.std(X, axis=0, ddof=1) + n_samples, n_dimensions = X.shape + normal_reference_factor = (4 / (n_dimensions + 2)) ** (1 / (n_dimensions + 4)) + scott_factor = normal_reference_factor * n_samples ** (-1 / (n_dimensions + 4)) + isotropic_scale = np.prod(scales) ** (1 / n_dimensions) + return float(scott_factor * isotropic_scale) + + def coefficient_of_determination( X: jax.Array, Y: jax.Array, diff --git a/tests/test_environment.py b/tests/test_environment.py index 51e489c..038418f 100644 --- a/tests/test_environment.py +++ b/tests/test_environment.py @@ -38,7 +38,7 @@ def test_coords(self): class TestEnvironment3D: def test_correct_dimensions(self): X = np.random.randn(100, 3) - env = Environment(X) + env = Environment(X, bin_size="auto") assert env.D == 3 assert env.dim == ["x", "y", "z"] @@ -76,7 +76,7 @@ class TestEnvironmentForceLims: def test_overrides_data_limits(self): X = np.random.randn(100, 2) lims = ((-5.0, -5.0), (5.0, 5.0)) - env = Environment(X, force_lims=lims) + env = Environment(X, force_lims=lims, bin_size=1.0) assert env.lims == lims diff --git a/tests/test_simpl.py b/tests/test_simpl.py index 00ab640..121a347 100644 --- a/tests/test_simpl.py +++ b/tests/test_simpl.py @@ -7,7 +7,7 @@ import xarray as xr from simpl.simpl import SIMPL -from simpl.utils import load_results +from simpl.utils import _estimate_kernel_bandwidth, load_results class TestSIMPLInit: @@ -112,6 +112,30 @@ def test_fit_sets_attributes(self, small_simpl_model): assert model.D_ == 2 assert model.iteration_ == 0 + def test_auto_parameters_are_inferred_saved_and_printed(self, demo_data, capsys): + N = 500 + Xb = demo_data["Xb"][:N] + time = demo_data["time"][:N] + model = SIMPL(speckle_block_size_seconds=0.1) + model.fit(demo_data["Y"][:N, :5], Xb, time, n_iterations=0) + + assert model.kernel_bandwidth == "auto" + expected_bandwidth = max(_estimate_kernel_bandwidth(Xb), model.bin_size_) + assert model.kernel_bandwidth_ == pytest.approx(expected_bandwidth) + assert model.results_.attrs["kernel_bandwidth"] == pytest.approx(model.kernel_bandwidth_) + + expected_speed = np.mean(np.linalg.norm(np.diff(Xb, axis=0), axis=1) / np.diff(time)) + assert model.speed_prior == "auto" + assert model.speed_prior_ == pytest.approx(expected_speed) + assert model.speed_prior_effective_ == pytest.approx(expected_speed) + assert model.results_.attrs["speed_prior"] == pytest.approx(expected_speed) + + output = capsys.readouterr().out + assert f"kernel_bandwidth={model.kernel_bandwidth_:.3f} (auto)" in output + assert f"bin_size={model.bin_size_:.3f} (auto)" in output + assert f"speed_prior={model.speed_prior_:.3f} (auto)" in output + assert "behavior_prior=None" in output + def test_fit_creates_environment_internally(self, demo_data): N = 500 N_neurons = min(5, demo_data["Y"].shape[1]) @@ -651,6 +675,19 @@ def test_align_angular_trajectory_mode(self): class TestSIMPLCircularEnvironment: + def test_auto_speed_prior_wraps_boundary_displacements(self): + T = 100 + time = np.arange(T) * 0.02 + unwrapped = np.linspace(np.pi - 0.2, np.pi + 0.2, T) + Xb = ((unwrapped + np.pi) % (2 * np.pi) - np.pi)[:, None] + model = SIMPL(is_1D_angular=True, kernel_bandwidth=0.3, speckle_block_size_seconds=0.1) + + model.fit(np.zeros((T, 3)), Xb, time, n_iterations=0, verbose=False) + + wrapped_displacement = (np.diff(Xb, axis=0) + np.pi) % (2 * np.pi) - np.pi + expected = np.mean(np.linalg.norm(wrapped_displacement, axis=1) / np.diff(time)) + assert model.speed_prior_ == pytest.approx(expected) + def test_circular_fit_uses_full_domain_even_with_partial_behavior(self): T, N_neurons = 200, 4 time = np.arange(T) * 0.02 diff --git a/tests/test_utils.py b/tests/test_utils.py index 49ceaf5..07efff5 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -11,6 +11,7 @@ _bin_indices_minuspi_pi, _circular_conv_fft_1d, _circular_mean_and_variance, + _estimate_kernel_bandwidth, _wrap_minuspi_pi, accumulate_spikes, analyse_place_fields, @@ -33,6 +34,19 @@ ) +class TestEstimateKernelBandwidth: + def test_uses_multivariate_scott_rule(self): + X = np.column_stack([np.linspace(0, 1, 100), np.linspace(-2, 2, 100)]) + scales = np.std(X, axis=0, ddof=1) + scott_factor = (4 / (2 + 2)) ** (1 / (2 + 4)) * len(X) ** (-1 / (2 + 4)) + expected = scott_factor * np.sqrt(np.prod(scales)) + + assert _estimate_kernel_bandwidth(X) == pytest.approx(expected) + + def test_zero_variance_returns_zero(self): + assert _estimate_kernel_bandwidth(np.ones((10, 2))) == 0.0 + + @pytest.mark.cpu_only class TestGaussianPdf: def test_correct_shape(self):