perturbvi 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,63 @@
1
+ # version is auto-built by hatch, not need to store it...
2
+ src/perturbvi/_version.py
3
+
4
+ test?.py
5
+
6
+ *.cache
7
+ uv.lock
8
+
9
+ # Temporary and binary files
10
+ foo.py
11
+ *~
12
+ *.py[cod]
13
+ *.so
14
+ *.cfg
15
+ !.isort.cfg
16
+ !setup.cfg
17
+ *.orig
18
+ *.log
19
+ *.pot
20
+ __pycache__/*
21
+ .cache/*
22
+ .*.swp
23
+ */.ipynb_checkpoints/*
24
+ .DS_Store
25
+
26
+ # Project files
27
+ .ropeproject
28
+ .project
29
+ .pydevproject
30
+ .settings
31
+ .idea
32
+ .vscode
33
+ tags
34
+
35
+ # Package files
36
+ *.egg
37
+ *.eggs/
38
+ .installed.cfg
39
+ *.egg-info
40
+
41
+ # Unittest and coverage
42
+ htmlcov/*
43
+ .coverage
44
+ .coverage.*
45
+ .tox
46
+ junit*.xml
47
+ coverage.xml
48
+ .pytest_cache/
49
+
50
+ # Build and docs folder/files
51
+ build/*
52
+ dist/*
53
+ sdist/*
54
+ docs/_rst/*
55
+ docs/_build/*
56
+ cover/*
57
+ site/
58
+ MANIFEST
59
+
60
+ # Per-project virtualenvs
61
+ .venv*/
62
+ .conda*/
63
+ .python-version
@@ -0,0 +1,9 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2023-present Nicholas Mancuso <nmancuso@usc.edu>
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
6
+
7
+ The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
8
+
9
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
@@ -0,0 +1,140 @@
1
+ Metadata-Version: 2.4
2
+ Name: perturbvi
3
+ Version: 0.1.0
4
+ Summary: Infer regulatory modules through informative latent component model in the single-cell Perturb-seq data
5
+ Project-URL: Documentation, https://mancusolab.github.io/perturbvi
6
+ Project-URL: Issues, https://github.com/mancusolab/perturbvi/issues
7
+ Project-URL: Source, https://github.com/mancusolab/perturbvi
8
+ Author-email: Dong Yuan <dongyuan@usc.edu>, Nicholas Mancuso <nmancuso@usc.edu>
9
+ License-Expression: MIT
10
+ License-File: LICENSE.txt
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Programming Language :: Python
13
+ Classifier: Programming Language :: Python :: 3.9
14
+ Classifier: Programming Language :: Python :: 3.10
15
+ Classifier: Programming Language :: Python :: 3.11
16
+ Classifier: Programming Language :: Python :: Implementation :: CPython
17
+ Classifier: Programming Language :: Python :: Implementation :: PyPy
18
+ Requires-Python: >=3.9
19
+ Requires-Dist: adjusttext
20
+ Requires-Dist: equinox
21
+ Requires-Dist: importlib-metadata; python_version < '3.8'
22
+ Requires-Dist: jax
23
+ Requires-Dist: jaxlib
24
+ Requires-Dist: jaxtyping
25
+ Requires-Dist: lineax
26
+ Requires-Dist: matplotlib
27
+ Requires-Dist: numpy
28
+ Requires-Dist: optax
29
+ Requires-Dist: optimistix
30
+ Requires-Dist: pandas
31
+ Requires-Dist: plum-dispatch
32
+ Requires-Dist: seaborn
33
+ Description-Content-Type: text/markdown
34
+
35
+ [![Documentation-webpage](https://img.shields.io/badge/Docs-Available-brightgreen)](https://mancusolab.github.io/perturbvi/)
36
+ [![PyPI-Server](https://img.shields.io/pypi/v/perturbvi.svg)](https://pypi.org/project/perturbvi/)
37
+ [![Github](https://img.shields.io/github/stars/mancusolab/perturbvi?style=social)](https://github.com/mancusolab/perturbvi)
38
+ [![License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
39
+ [![Project generated with Hatch](https://img.shields.io/badge/%F0%9F%A5%9A-Hatch-4051b5.svg)](https://github.com/pypa/hatch)
40
+
41
+ # perturbVI
42
+ `perturbvi` is a scalable approach to infer regulatory modules through informative latent component model in the single-cell Perturb-seq data.
43
+
44
+ [**Installation**](#installation)
45
+ | [**Example**](#get-started-with-example)
46
+ | [**Notes**](#notes)
47
+ | [**Version**](#version-history)
48
+ | [**Support**](#support)
49
+ | [**Other Software**](#other-software)
50
+
51
+ ------------------
52
+
53
+ ## Installation
54
+
55
+ Users can download the latest repository and then use `pip`:
56
+
57
+ ``` bash
58
+ git clone https://github.com/mancusolab/perturbvi.git
59
+ cd perturbvi
60
+ pip install .
61
+ ```
62
+
63
+ ## Get Started with `perturbvi`
64
+
65
+ ### 1. `infer`
66
+ Perform inference using SuSiE PCA to find the regulatory modules from CRISPR perturbation data
67
+ ``` bash
68
+ perturbvi infer <exp_csv> <guide_csv> <gene_symbol_csv> -o=output --verbose
69
+ ```
70
+
71
+ #### Arguments
72
+ - `exp_csv`: Path to the experiment CSV file.
73
+ - `guide_csv`: Path to the guide CSV file.
74
+ - `gene_symbol_csv`: Path to the gene symbol CSV file.
75
+ - `-o=output`: Specifies the output directory name or path.
76
+ - `--verbose`: For logging (Optional).
77
+
78
+ #### Example Usage
79
+ ```bash
80
+ perturbvi infer data/exp.csv data/guide.csv data/symbol.csv -o=data/out --verbose
81
+ ```
82
+
83
+ This will save the all the output files (including the parameter file `params.pkl`) into the `data/out` folder, which can be used for the downstream tasks outlined below.
84
+
85
+ ## Notes
86
+
87
+ - `perturbvi` uses [JAX](https://github.com/google/jax) with [Just In
88
+ Time](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html)
89
+ compilation to achieve high-speed computation. However, there are
90
+ some [issues](https://github.com/google/jax/issues/5501) for JAX
91
+ with Mac M1 chip. To solve this, users need to initiate conda using
92
+ [miniforge](https://github.com/conda-forge/miniforge), and then
93
+ install `perturbvi` using `pip` in the desired environment.
94
+
95
+ ## Version History
96
+
97
+ TBD
98
+
99
+ ## Support
100
+
101
+ Please report any bugs or feature requests in the [Issue
102
+ Tracker](https://github.com/mancusolab/perturbvi/issues). If users have
103
+ any questions or comments, please contact Dong Yuan (<dongyuan@usc.edu>)
104
+ and Nicholas Mancuso (<nmancuso@usc.edu>).
105
+
106
+ ## Other Software
107
+
108
+ Feel free to use other software developed by [Mancuso
109
+ Lab](https://www.mancusolab.com/):
110
+
111
+ - [SuShiE](https://github.com/mancusolab/sushie): a Bayesian
112
+ fine-mapping framework for molecular QTL data across multiple
113
+ ancestries.
114
+ - [MA-FOCUS](https://github.com/mancusolab/ma-focus): a Bayesian
115
+ fine-mapping framework using
116
+ [TWAS](https://www.nature.com/articles/ng.3506) statistics across
117
+ multiple ancestries to identify the causal genes for complex traits.
118
+ - [SuSiE-PCA](https://github.com/mancusolab/susiepca): a scalable
119
+ Bayesian variable selection technique for sparse principal component
120
+ analysis
121
+ - [twas_sim](https://github.com/mancusolab/twas_sim): a Python
122
+ software to simulate [TWAS](https://www.nature.com/articles/ng.3506)
123
+ statistics.
124
+ - [FactorGo](https://github.com/mancusolab/factorgo): a scalable
125
+ variational factor analysis model that learns pleiotropic factors
126
+ from GWAS summary statistics.
127
+ - [HAMSTA](https://github.com/tszfungc/hamsta): a Python software to
128
+ estimate heritability explained by local ancestry data from
129
+ admixture mapping summary statistics.
130
+
131
+ ------------------------------------------------------------------------
132
+
133
+ `perturbvi` is distributed under the terms of the
134
+ [MIT](https://spdx.org/licenses/MIT.html) license.
135
+
136
+
137
+ ------------------------------------------------------------------------
138
+
139
+ This project has been set up using Hatch. For details and usage
140
+ information on Hatch see <https://github.com/pypa/hatch>.
@@ -0,0 +1,106 @@
1
+ [![Documentation-webpage](https://img.shields.io/badge/Docs-Available-brightgreen)](https://mancusolab.github.io/perturbvi/)
2
+ [![PyPI-Server](https://img.shields.io/pypi/v/perturbvi.svg)](https://pypi.org/project/perturbvi/)
3
+ [![Github](https://img.shields.io/github/stars/mancusolab/perturbvi?style=social)](https://github.com/mancusolab/perturbvi)
4
+ [![License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
5
+ [![Project generated with Hatch](https://img.shields.io/badge/%F0%9F%A5%9A-Hatch-4051b5.svg)](https://github.com/pypa/hatch)
6
+
7
+ # perturbVI
8
+ `perturbvi` is a scalable approach to infer regulatory modules through informative latent component model in the single-cell Perturb-seq data.
9
+
10
+ [**Installation**](#installation)
11
+ | [**Example**](#get-started-with-example)
12
+ | [**Notes**](#notes)
13
+ | [**Version**](#version-history)
14
+ | [**Support**](#support)
15
+ | [**Other Software**](#other-software)
16
+
17
+ ------------------
18
+
19
+ ## Installation
20
+
21
+ Users can download the latest repository and then use `pip`:
22
+
23
+ ``` bash
24
+ git clone https://github.com/mancusolab/perturbvi.git
25
+ cd perturbvi
26
+ pip install .
27
+ ```
28
+
29
+ ## Get Started with `perturbvi`
30
+
31
+ ### 1. `infer`
32
+ Perform inference using SuSiE PCA to find the regulatory modules from CRISPR perturbation data
33
+ ``` bash
34
+ perturbvi infer <exp_csv> <guide_csv> <gene_symbol_csv> -o=output --verbose
35
+ ```
36
+
37
+ #### Arguments
38
+ - `exp_csv`: Path to the experiment CSV file.
39
+ - `guide_csv`: Path to the guide CSV file.
40
+ - `gene_symbol_csv`: Path to the gene symbol CSV file.
41
+ - `-o=output`: Specifies the output directory name or path.
42
+ - `--verbose`: For logging (Optional).
43
+
44
+ #### Example Usage
45
+ ```bash
46
+ perturbvi infer data/exp.csv data/guide.csv data/symbol.csv -o=data/out --verbose
47
+ ```
48
+
49
+ This will save the all the output files (including the parameter file `params.pkl`) into the `data/out` folder, which can be used for the downstream tasks outlined below.
50
+
51
+ ## Notes
52
+
53
+ - `perturbvi` uses [JAX](https://github.com/google/jax) with [Just In
54
+ Time](https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html)
55
+ compilation to achieve high-speed computation. However, there are
56
+ some [issues](https://github.com/google/jax/issues/5501) for JAX
57
+ with Mac M1 chip. To solve this, users need to initiate conda using
58
+ [miniforge](https://github.com/conda-forge/miniforge), and then
59
+ install `perturbvi` using `pip` in the desired environment.
60
+
61
+ ## Version History
62
+
63
+ TBD
64
+
65
+ ## Support
66
+
67
+ Please report any bugs or feature requests in the [Issue
68
+ Tracker](https://github.com/mancusolab/perturbvi/issues). If users have
69
+ any questions or comments, please contact Dong Yuan (<dongyuan@usc.edu>)
70
+ and Nicholas Mancuso (<nmancuso@usc.edu>).
71
+
72
+ ## Other Software
73
+
74
+ Feel free to use other software developed by [Mancuso
75
+ Lab](https://www.mancusolab.com/):
76
+
77
+ - [SuShiE](https://github.com/mancusolab/sushie): a Bayesian
78
+ fine-mapping framework for molecular QTL data across multiple
79
+ ancestries.
80
+ - [MA-FOCUS](https://github.com/mancusolab/ma-focus): a Bayesian
81
+ fine-mapping framework using
82
+ [TWAS](https://www.nature.com/articles/ng.3506) statistics across
83
+ multiple ancestries to identify the causal genes for complex traits.
84
+ - [SuSiE-PCA](https://github.com/mancusolab/susiepca): a scalable
85
+ Bayesian variable selection technique for sparse principal component
86
+ analysis
87
+ - [twas_sim](https://github.com/mancusolab/twas_sim): a Python
88
+ software to simulate [TWAS](https://www.nature.com/articles/ng.3506)
89
+ statistics.
90
+ - [FactorGo](https://github.com/mancusolab/factorgo): a scalable
91
+ variational factor analysis model that learns pleiotropic factors
92
+ from GWAS summary statistics.
93
+ - [HAMSTA](https://github.com/tszfungc/hamsta): a Python software to
94
+ estimate heritability explained by local ancestry data from
95
+ admixture mapping summary statistics.
96
+
97
+ ------------------------------------------------------------------------
98
+
99
+ `perturbvi` is distributed under the terms of the
100
+ [MIT](https://spdx.org/licenses/MIT.html) license.
101
+
102
+
103
+ ------------------------------------------------------------------------
104
+
105
+ This project has been set up using Hatch. For details and usage
106
+ information on Hatch see <https://github.com/pypa/hatch>.
@@ -0,0 +1,144 @@
1
+ [build-system]
2
+ requires = ["hatchling", "hatch-vcs"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "perturbvi"
7
+ version="0.1.0"
8
+ description = 'Infer regulatory modules through informative latent component model in the single-cell Perturb-seq data'
9
+ readme = "README.md"
10
+ requires-python = ">=3.9"
11
+ license = "MIT"
12
+ keywords = []
13
+ authors = [
14
+ { name = "Dong Yuan", email = "dongyuan@usc.edu" },
15
+ { name = "Nicholas Mancuso", email = "nmancuso@usc.edu" },
16
+ ]
17
+ classifiers = [
18
+ "Development Status :: 4 - Beta",
19
+ "Programming Language :: Python",
20
+ "Programming Language :: Python :: 3.9",
21
+ "Programming Language :: Python :: 3.10",
22
+ "Programming Language :: Python :: 3.11",
23
+ "Programming Language :: Python :: Implementation :: CPython",
24
+ "Programming Language :: Python :: Implementation :: PyPy",
25
+ ]
26
+ dependencies = [
27
+ "importlib-metadata; python_version<\"3.8\"",
28
+ "jaxlib",
29
+ "jax",
30
+ "jaxtyping",
31
+ "equinox",
32
+ "lineax",
33
+ "optimistix",
34
+ "optax",
35
+ "plum-dispatch",
36
+ "pandas",
37
+ "numpy",
38
+ "matplotlib",
39
+ "seaborn",
40
+ "adjustText"
41
+ ]
42
+
43
+ [project.scripts]
44
+ perturbvi = "perturbvi.cli:run_cli"
45
+
46
+ [project.urls]
47
+ Documentation = "https://mancusolab.github.io/perturbvi"
48
+ Issues = "https://github.com/mancusolab/perturbvi/issues"
49
+ Source = "https://github.com/mancusolab/perturbvi"
50
+
51
+ [tool.hatch.version]
52
+ source = "vcs"
53
+
54
+ [tool.hatch.build.hooks.vcs]
55
+ version-file = "src/perturbvi/_version.py"
56
+
57
+ [tool.hatch.build.targets.sdist]
58
+ include = ["src/"]
59
+ exclude = ["docs/"]
60
+
61
+ [tool.hatch.envs.test]
62
+ dependencies = [
63
+ "coverage[toml]",
64
+ "pytest",
65
+ "pytest-cov",
66
+ ]
67
+
68
+ [tool.hatch.envs.test.scripts]
69
+ run-coverage = "pytest --cov-config=pyproject.toml --cov=pkg --cov=tests"
70
+ run = "run-coverage --no-cov"
71
+
72
+ [[tool.hatch.envs.all.matrix]]
73
+ python = ["3.9", "3.10", "3.11"]
74
+
75
+ [tool.hatch.envs.lint]
76
+ detached = true
77
+ dependencies = [
78
+ "black>=23.1.0",
79
+ "mypy>=1.0.0",
80
+ "ruff>=0.0.243",
81
+ ]
82
+
83
+ [tool.ruff]
84
+ target-version = "py37"
85
+ line-length = 120
86
+
87
+ [tool.ruff.lint]
88
+ select = ["E", "F", "I001"]
89
+ ignore = [
90
+ # Allow non-abstract empty methods in abstract base classes
91
+ "B027",
92
+ # Allow boolean positional values in function calls, like `dict.get(... True)`
93
+ "FBT003",
94
+ # Ignore checks for possible passwords
95
+ "S105", "S106", "S107",
96
+ # Ignore complexity
97
+ "C901", "PLR0911", "PLR0912", "PLR0913", "PLR0915",
98
+ # Ignore syntax checks on jaxarray shaped typing
99
+ "F722",
100
+ # Ignore multidispatch false positives
101
+ "F811",
102
+ ]
103
+
104
+ [tool.ruff.lint.isort]
105
+ known-first-party = ["susieguy"]
106
+ combine-as-imports = true
107
+ lines-after-imports = 2
108
+ lines-between-types = 1
109
+ known-local-folder = ["src"]
110
+ section-order = ["future", "standard-library", "third-party", "jax-ecosystem", "first-party", "local-folder"]
111
+ extra-standard-library = ["typing_extensions"]
112
+ order-by-type = false
113
+
114
+ [tool.ruff.lint.isort.sections]
115
+ jax-ecosystem = ["equinox", "jax", "jaxtyping", "lineax", "optimistix", "optax"]
116
+
117
+ [tool.ruff.lint.per-file-ignores]
118
+ # Tests can use magic values, assertions, and relative imports
119
+ "tests/**/*" = ["PLR2004", "S101", "TID252"]
120
+
121
+ [tool.pyright]
122
+ reportIncompatibleMethodOverride = true
123
+ reportIncompatibleVariableOverride = false # Incompatible with eqx.AbstractVar
124
+ include = ["perturbvi", "tests"]
125
+ exclude = ["docs"]
126
+
127
+ [tool.coverage.run]
128
+ source_pkgs = ["perturbvi", "tests"]
129
+ branch = true
130
+ parallel = true
131
+ omit = [
132
+ "src/susieguy/__about__.py",
133
+ ]
134
+
135
+ [tool.coverage.paths]
136
+ susieguy = ["src/perturbvi", "*/perturbvi/src/perturbvi"]
137
+ tests = ["tests", "*/perturbvi/tests"]
138
+
139
+ [tool.coverage.report]
140
+ exclude_lines = [
141
+ "no cov",
142
+ "if __name__ == .__main__.:",
143
+ "if TYPE_CHECKING:",
144
+ ]
@@ -0,0 +1,22 @@
1
+ from importlib.metadata import PackageNotFoundError, version # pragma: no cover
2
+
3
+ # annoying lint bug...
4
+ from .infer import (
5
+ compute_elbo as compute_elbo,
6
+ compute_pip as compute_pip,
7
+ compute_pve as compute_pve,
8
+ infer as infer,
9
+ )
10
+ from .sim import generate_sim as generate_sim
11
+ from .utils import compute_lfsr as compute_lfsr
12
+ from .io import save_results as save_results
13
+
14
+
15
+ try:
16
+ # Change here if project is renamed and does not equal the package name
17
+ dist_name = __name__
18
+ __version__ = version(dist_name)
19
+ except PackageNotFoundError: # pragma: no cover
20
+ __version__ = "unknown"
21
+ finally:
22
+ del version, PackageNotFoundError
@@ -0,0 +1,21 @@
1
+ # file generated by setuptools-scm
2
+ # don't change, don't track in version control
3
+
4
+ __all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
5
+
6
+ TYPE_CHECKING = False
7
+ if TYPE_CHECKING:
8
+ from typing import Tuple
9
+ from typing import Union
10
+
11
+ VERSION_TUPLE = Tuple[Union[int, str], ...]
12
+ else:
13
+ VERSION_TUPLE = object
14
+
15
+ version: str
16
+ __version__: str
17
+ __version_tuple__: VERSION_TUPLE
18
+ version_tuple: VERSION_TUPLE
19
+
20
+ __version__ = version = '0.1.0'
21
+ __version_tuple__ = version_tuple = (0, 1, 0)
@@ -0,0 +1,80 @@
1
+ from abc import abstractmethod
2
+ from dataclasses import field
3
+ from typing import Any
4
+
5
+ import equinox as eqx
6
+ import jax
7
+ import jax.nn as nn
8
+ import optimistix as optx
9
+
10
+ from jaxtyping import Array
11
+
12
+ from .common import DataMatrix, ModelParams
13
+ from .utils import kl_discrete
14
+
15
+
16
+ def _compute_pi(A: DataMatrix, theta: Array) -> Array:
17
+ return nn.softmax(A @ theta, axis=0).T
18
+
19
+
20
+ def _loss(theta: Array, args) -> Array:
21
+ A, alpha = args
22
+ pi = _compute_pi(A, theta)
23
+ return kl_discrete(alpha, pi), None
24
+
25
+
26
+ class PriorModel(eqx.Module):
27
+ @abstractmethod
28
+ def predict(self, params: ModelParams) -> Array: ...
29
+
30
+ @abstractmethod
31
+ def init_state(self, params: ModelParams) -> ModelParams: ...
32
+
33
+ @abstractmethod
34
+ def update(self, params: ModelParams) -> ModelParams: ...
35
+
36
+
37
+ class FixedPrior(PriorModel):
38
+ def predict(self, params: ModelParams) -> Array:
39
+ return params.pi
40
+
41
+ def init_state(self, params: ModelParams) -> ModelParams:
42
+ return params
43
+
44
+ def update(self, params: ModelParams) -> ModelParams:
45
+ return params
46
+
47
+
48
+ class AnnotationPriorModel(PriorModel):
49
+ A: DataMatrix
50
+ search: optx.AbstractMinimiser
51
+ step: Any = field(init=False)
52
+
53
+ def __post_init__(self):
54
+ self.step = eqx.filter_jit(eqx.Partial(self.search.step, _loss, options=None, tags=None))
55
+
56
+ @property
57
+ def shape(self):
58
+ return self.A.shape
59
+
60
+ def init_state(self, params: ModelParams) -> ModelParams:
61
+ args = (self.A, params.alpha)
62
+ f_struct, aux_struct = jax.eval_shape(_loss, params.theta, args)
63
+ return params._replace(
64
+ ann_state=self.search.init(
65
+ _loss, params.theta, args, options=None, f_struct=f_struct, aux_struct=aux_struct, tags=None
66
+ ),
67
+ )
68
+
69
+ def predict(self, params: ModelParams) -> Array:
70
+ return _compute_pi(self.A, params.theta)
71
+
72
+ def update(self, params: ModelParams) -> ModelParams:
73
+ args = (self.A, params.alpha)
74
+ # take one step using optimistix optimizer
75
+ theta, state, _ = self.step(params.theta, state=params.ann_state, args=args)
76
+ return params._replace(
77
+ theta=theta,
78
+ pi=_compute_pi(self.A, theta),
79
+ ann_state=state,
80
+ )