SeInE-orientation 0.1.0__py3-none-any.whl

seine_orientation/BayesModel.py

@@ -0,0 +1,283 @@
+import numpy as np
+import time
+import itertools
+
+from seine.NestedSamplingMethods import (
+    run_sampling,
+)
+from seine import HierarchicalModel, functions as prior_fn
+from seine.structures import (
+    prior_structure,
+)
+
+from .utils import gabor_filter, gabor_response, sine_grating, softplus, ReLU, sigmoid
+
+
+class HierarchicalBayesInference(HierarchicalModel):
+
+    def set_gratings(self, measure_points, FoV_range, FoV_steps):
+
+        self.measure_points = measure_points
+
+        self.X_FoV, self.Y_FoV = np.meshgrid(
+            np.linspace(-FoV_range[0] / 2, FoV_range[0] / 2, FoV_steps),
+            np.linspace(-FoV_range[1] / 2, FoV_range[1] / 2, FoV_steps),
+        )
+        self.FoV_grid = np.dstack((self.X_FoV, self.Y_FoV))
+
+        FoV_steps = self.X_FoV.shape[0]
+
+        self.gratings = np.zeros(
+            (
+                *[len(mp) for mp in self.measure_points],
+                FoV_steps,
+                FoV_steps,
+            )
+        )
+        for prod in itertools.product(
+            enumerate(self.measure_points[0]),
+            enumerate(self.measure_points[1]),
+            enumerate(self.measure_points[2]),
+        ):
+            idx, elems = zip(*prod)
+            phi_0, theta, f = elems
+            self.gratings[*idx, ...] = sine_grating(
+                self.X_FoV,  # [0, ...],
+                self.Y_FoV,  # [0, ...],
+                theta,
+                f,
+                phi_0,
+                square=True,
+            )
+
+    def set_priors(self, priors_init=None, coding="simple"):
+
+        if priors_init is None:
+
+            if hasattr(self, "data"):
+                fmap = self.data["observed_counts"] / self.data["T"]
+                A0_guess, A_guess = np.percentile(fmap, [50, 90])
+                A0_guess = np.maximum(A0_guess, 1.0)
+                A_guess = np.maximum(A_guess, 2.0)
+            else:
+                A0_guess, A_guess = 1.0, 3.0
+            self.priors_init = {}
+            # print("A guesses from data:", A0_guess, A_guess)
+
+            ## define gabor model priors
+            if coding in ["simple", "complex"]:
+                self.priors_init["theta"] = prior_structure(
+                    prior_fn.bounded_flat,
+                    low=-np.pi / 2,
+                    high=np.pi / 2,
+                    label=r"$\theta$",
+                    periodic=True,
+                )
+                self.priors_init["f"] = prior_structure(
+                    prior_fn.halfnorm_ppf,
+                    loc=0.0,
+                    scale=5.0,
+                    label=r"$f$",
+                )
+                if coding == "simple":
+                    ## not required for complex cell model
+                    self.priors_init["phi_0"] = prior_structure(
+                        prior_fn.bounded_flat,
+                        low=-np.pi,
+                        high=np.pi,
+                        periodic=True,
+                        label=r"$\phi_0$",
+                    )
+                else:
+                    self.priors_init["phi_0"] = prior_structure(
+                        None,
+                        value=0.0,
+                        periodic=True,
+                        label=r"$\phi_0$",
+                    )
+
+                self.priors_init["sigma"] = prior_structure(
+                    prior_fn.halfnorm_ppf,
+                    loc=0.05,
+                    scale=0.1,
+                    label=r"$\sigma$",
+                )
+                self.priors_init["gamma"] = prior_structure(
+                    prior_fn.halfnorm_ppf,
+                    loc=1.0,
+                    scale=1.0,
+                    label=r"$\gamma$",
+                )
+                self.priors_init["theta_gauss"] = prior_structure(
+                    prior_fn.bounded_flat,
+                    low=-np.pi / 2,
+                    high=np.pi / 2,
+                    periodic=True,
+                    label=r"$\theta_{gauss}$",
+                )
+
+            ## define nonlinearity priors - should start with "nl_" to be recognized by model response function
+            self.priors_init["nl_baseline"] = prior_structure(
+                prior_fn.halfnorm_ppf,
+                loc=0.0,
+                scale=A0_guess,
+                label=r"$A_{baseline}$",
+            )
+            if coding in ["simple", "complex"]:
+                self.priors_init["nl_transition"] = prior_structure(
+                    prior_fn.halfnorm_ppf,
+                    loc=0.0,
+                    scale=1.0,
+                    label=r"$A_{transition}$",
+                )
+                self.priors_init["nl_amplitude"] = prior_structure(
+                    prior_fn.halfnorm_ppf,
+                    loc=A0_guess * 0.2,
+                    scale=A_guess - A0_guess,
+                    label=r"$A_{amplitude}$",
+                )
+
+            # parameters for loglikelihood
+            self.priors_init["logl_alpha"] = prior_structure(
+                prior_fn.halfnorm_ppf,
+                loc=0.0,
+                scale=1.0,
+                label=r"$\alpha_{logl}$",
+            )
+
+        else:
+            self.priors_init = priors_init
+
+        super().set_priors(self.priors_init)
+
+    def get_model_rate_response(self, params, coding="simple"):
+
+        ## transform parameters, if not already provided in proper format
+        if not isinstance(params, dict):
+            params = self.get_params_from_p(params)
+
+        ## get (normalized) model response
+        if coding == "random":
+            response = np.zeros(self.dimensions["shape"])
+        else:
+            response = gabor_response(
+                self.X_FoV,
+                self.Y_FoV,
+                params,
+                self.gratings,
+                mode=coding,
+                significance_threshold=0.01,
+                logger=self.log,
+            )
+        self.timeit("calculating model and response firing rate")
+
+        ## apply nonlinearity to adjust to scale of firing rates and get final model response in firing rate units
+        rate_response = self.apply_nonlinearity(response, params)
+        return rate_response
+
+    def apply_nonlinearity(self, response, params, nonlinearity="softplus"):
+
+        if nonlinearity == "softplus":
+            # print(
+            #     f"Applying softplus nonlinearity with parameters: baseline={params.get('nl_baseline', 0.0):.3f}, transition={params.get('nl_transition', 1.0):.3f}, amplitude={params.get('nl_amplitude', 1.0):.3f}"
+            # )
+            rate_response = softplus(
+                response * params.get("nl_amplitude", 1.0),
+                alpha=params.get("nl_transition", 1.0),
+                delta=params.get("nl_baseline", 0.0),
+            )
+        elif nonlinearity == "ReLU":
+            # rate_response =
+            raise NotImplementedError("ReLU nonlinearity not implemented yet")
+        elif nonlinearity == "sigmoid":
+            raise NotImplementedError("Sigmoid nonlinearity not implemented yet")
+        else:
+            raise ValueError(f"Unknown nonlinearity: {nonlinearity}")
+
+        self.timeit("applying nonlinearity")
+        return rate_response
+
+    def set_logp_func(self, coding="simple"):
+        """
+        some nice description
+        """
+
+        def get_logp(p_in):
+
+            self.timeit()
+            """
+                build switch between 3 models:
+                    - gabor type model
+                    - fourier component model (with n components in each direction)
+                    - ellipse in rate space
+            """
+
+            params = self.get_params_from_p(p_in)
+            self.timeit("transforming parameters")
+
+            model_rate_response = self.get_model_rate_response(params, coding)
+            if not (model_rate_response.shape == self.data["observed_counts"].shape):
+                model_rate_response = model_rate_response[None, ...]
+
+            logp = self.probability_of_spike_observation(
+                model_rate_response,
+                model="negative_binomial",
+                **params,
+            )
+            self.timeit("calculating log probability of spike observation")
+
+            return logp.sum()
+
+        return get_logp
+
+
+def run_inference(
+    hbm: HierarchicalBayesInference,
+    coding="simple",
+    n_live=200,
+    nP=1,
+    dlogz=1.0,
+    show_status=True,
+):
+    hbm.set_priors(coding=coding)
+    my_trafo = hbm.set_prior_transform()
+    my_logp = hbm.set_logp_func(coding=coding)
+
+    results, sampler = run_sampling(
+        my_trafo,
+        my_logp,
+        hbm.parameter_names_all,
+        hbm.periodic,
+        show_status=show_status,
+        n_live=n_live,
+        nP=nP,
+        dlogz=dlogz,
+    )
+    return results
+
+
+def run_model_comparison(data, dwelltime, measure_points, show_status=True, **kwargs):
+
+    t_start = time.time()
+    hbm = HierarchicalBayesInference()
+
+    hbm.set_gratings(
+        measure_points,
+        [np.deg2rad(140), np.deg2rad(114)],
+        FoV_steps=51,
+    )
+
+    hbm.prepare_data(
+        data,
+        dwelltime,
+        iter_dims=False,
+    )
+
+    results = {}
+    for coding in ["random", "simple", "complex"]:
+        results[coding] = run_inference(
+            hbm, coding=coding, show_status=show_status, **kwargs
+        )
+    t_end = time.time()
+    print(f"Model comparison done after {t_end - t_start:.2f} seconds")
+    return results

seine_orientation/__init__.py

@@ -0,0 +1,2 @@
+from .BayesModel import HierarchicalBayesInference, run_inference, run_model_comparison
+from .process_session import process_session

seine_orientation/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = 'gffe43a317'

seine_orientation/process_session.py

@@ -0,0 +1,269 @@
+from __future__ import annotations
+
+import h5py, os
+import numpy as np
+from pathlib import Path
+import multiprocessing as mp
+from functools import partial
+from scipy.io import loadmat
+
+from pathlib import Path
+from typing import Optional
+
+import argparse
+
+
+def get_folder_path(
+    mouse,
+    session,
+    *,
+    type="data",
+    dataset="PSD95Mice_Loewel",
+    project="tmp",
+):
+    return Path(project, type, dataset, mouse, session)
+
+
+def process_session(
+    # define the path towards data, assuming hpc structure
+    animal: str,
+    session: str,
+    fname: str,
+    dataset: str = "PSD95Mice_Loewel",
+    project: str = "orientation_selectivity",
+    spikes_key: str = "/estimates/S_dff",
+    response_onset: float = 0.2,
+    prefix: str = "SeInE",
+    suffix: str = "",
+    nP: int = 12,
+    save_type: str = "hdf5",
+    **kwargs,
+):
+    """
+    this assumes some specific structure and naming of data,
+    following the HPC project directory conventions.
+    """
+
+    from .utils.utils_analysis import (
+        calculate_firing_maps,
+        get_spikes,
+        get_unique_stimulus_values,
+    )
+    from .BayesModel import (
+        run_model_comparison,
+    )
+
+    projects_dir = os.environ.get("PROJECT_DIR", "tmp")
+    project_dir = Path(projects_dir) / project
+
+    dir_analysis = get_folder_path(
+        animal,
+        session,
+        type="analysis",
+        dataset=dataset,
+        project=str(project_dir),
+    )
+    path_meta = dir_analysis / "CaimanMeta.mat"
+    print(path_meta)
+    path_detection = dir_analysis / fname
+    print(path_detection)
+
+    dir_data = get_folder_path(
+        animal,
+        session,
+        type="data",
+        dataset=dataset,
+        project=str(project_dir),
+    )
+
+    ## from here, process data
+    ld = loadmat(path_meta, variable_names=["CaimanMeta"])
+    meta_data = ld["CaimanMeta"]
+
+    with h5py.File(path_detection, "r") as f:
+        S = np.array(f[spikes_key][()])
+
+    protocols = {"bino": 0, "cont": 1, "ipsi": 2}
+
+    for protocol, idx in protocols.items():
+        idx = protocols[protocol]
+        num_frames = np.cumsum(meta_data["num_frames"])
+        start_idx = num_frames[idx - 1] if idx > 0 else 0
+        end_idx = num_frames[idx]
+
+        stimuli = meta_data["Stimulus"][idx]
+        S_protocol = S[:, start_idx:end_idx]
+
+        path_stimulus = dir_data.glob(f"*_{protocol}_*").__next__()
+        ld = loadmat(path_stimulus)
+        stimulus_data = ld["runInfo"]
+
+        unique_values = get_unique_stimulus_values(stimulus_data)
+        measure_points = (
+            np.deg2rad(unique_values["phases"]),
+            np.deg2rad(unique_values["angles"]),
+            1.0 / np.deg2rad(1.0 / unique_values["cycles"]),
+        )
+
+        f = meta_data["frame_rate"]
+        spikes = get_spikes(S_protocol, f=f)
+
+        ## calculate baseline rates before / after stimulus presentation
+        baseline_pre = spikes[:, : int(stimuli[0, 0])].mean(axis=1) * f
+        baseline_post = spikes[:, int(stimuli[-1, 1]) :].mean(axis=1) * f
+
+        ## calculate firing maps during stimulus presentation
+        event_counts, dwelltime = calculate_firing_maps(
+            stimulus_data=stimulus_data,
+            spikes=spikes,
+            f=f,
+            dt_onset=response_onset,
+            dt_offset=response_onset / 2,
+            stimulus_frames=stimuli[:, :4],
+            collapse_repeats=True,
+        )
+
+        neurons = range(event_counts.shape[-1])
+        idx_process = np.array(neurons)
+        n_neurons = len(neurons)
+
+        process_neuron = partial(
+            run_model_comparison,
+            dwelltime=dwelltime,
+            measure_points=measure_points,
+            show_status=False,
+        )
+        # fmaps = (event_counts / dwelltime[..., np.newaxis]).transpose(3, 0, 1, 2)
+        fmaps = event_counts.transpose(3, 0, 1, 2)
+
+        batch_sz = 10 * nP
+        nBatch = n_neurons // batch_sz
+
+        print(
+            f"Processing {n_neurons} neurons in {nBatch+1} batches using {nP} processes..."
+        )
+        results = []
+        with mp.Pool(nP) as pool:
+            for i in range(nBatch + 1):
+                idx_batch = idx_process[
+                    i * batch_sz : min(n_neurons, (i + 1) * batch_sz)
+                ]
+                outputs = pool.map(
+                    process_neuron,
+                    fmaps[idx_batch, ...],
+                )
+
+                for n, entry in zip(idx_batch, outputs):
+                    results.append({})
+                    for model in entry.keys():
+                        ## store only relevant output
+                        results[n][model] = {
+                            "fmap": fmaps[n, ...],
+                            "baseline": [baseline_pre[n], baseline_post[n]],
+                            "dwelltime": dwelltime,  ## is saved n_neuron times, but is not large
+                            "evidence": [
+                                entry[n][model].logz[-1],
+                                entry[n][model].logzerr[-1],
+                            ],
+                            "samples": entry[n][model].samples,
+                            "weights": entry[n][model].importance_weights(),
+                        }
+
+        fname_out = (
+            dir_analysis / f"{prefix}_{protocol}_{Path(fname).stem}{suffix}.{save_type}"
+        )
+        with h5py.File(fname_out, "w") as f:
+            for n, entry in enumerate(results):
+                grp_neuron = f.create_group(f"neuron_{n}")
+                for model in entry.keys():
+                    grp_model = grp_neuron.create_group(model)
+                    for key, value in entry[model].items():
+                        if isinstance(value, list):
+                            grp_model.create_dataset(key, data=np.array(value))
+                        else:
+                            grp_model.create_dataset(key, data=value)
+        print(f"Results from {n_neurons} neurons saved to {fname_out}")
+
+
+def build_parser() -> argparse.ArgumentParser:
+
+    parser = argparse.ArgumentParser(
+        prog="neuron-detection", description="Neuron detection pipeline using CaImAn"
+    )
+    parser.add_argument(
+        "--animal",
+        type=str,
+        required=True,
+        help="Identifier for the animal being processed",
+    )
+    parser.add_argument(
+        "--session",
+        type=str,
+        required=True,
+        help="Identifier for the session being processed",
+    )
+    parser.add_argument(
+        "--fname",
+        type=str,
+        required=True,
+        help="Filename of the imaging data file",
+    )
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="PSD95Mice_Loewel",
+        help="Name of the dataset being processed",
+    )
+    parser.add_argument(
+        "--project",
+        type=str,
+        default="orientation_selectivity",
+        help="Name of the project",
+    )
+    parser.add_argument(
+        "--spikes_key",
+        type=str,
+        default="/estimates/S_dff",
+        help="HDF5 key for the spikes data",
+    )
+    parser.add_argument(
+        "--response_onset",
+        type=float,
+        default=0.2,
+        help="Response onset time",
+    )
+
+    parser.add_argument(
+        "--prefix",
+        type=str,
+        default="results_CaImAn",
+        help="Prefix of CaImAn result files",
+    )
+    parser.add_argument(
+        "--suffix",
+        type=str,
+        default="",
+        help="Optional suffix for different runs of detection",
+    )
+    parser.add_argument(
+        "--nP",
+        type=int,
+        default=12,
+        help="Number of processes for parallel processing",
+    )
+    parser.add_argument(
+        "--save_type",
+        type=str,
+        default="hdf5",
+        help="Specifies result file type (without trailing '.'). Defaults to hdf5",
+    )
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    process_session(**args.__dict__)
+
+    return 0

seine_orientation/utils/__init__.py

@@ -0,0 +1,11 @@
+from .utils_display import *
+from .utils_analysis import *
+from .utils_model import (
+    gabor_filter,
+    gabor_response,
+    sine_grating,
+    softplus,
+    ReLU,
+    sigmoid,
+    gompertz,
+)