moseq2-model 1.4.0__tar.gz

moseq2_model-1.4.0/LICENSE.md

@@ -0,0 +1,49 @@
+# Non-Commercial Research and Academic Use Software License and Terms of Use
+
+
+MoSeq is a software package that includes original code created by the Harvard researchers listed below (the “Software”), and third-party code that may be obtained by End Users separately. The Software is designed to label behaviors identified from 3D pose dynamics data through the use of computational modeling and fitting approaches. The Software was developed by Alexander Wiltschko, Ph.D., Sandeep Robert Datta, Ph.D., and Matthew J. Johnson, Ph.D. at Harvard University. It is distributed for free academic and non-commercial research use by the President and Fellows of Harvard College (“Harvard”).
+
+Using the Software indicates your agreement to be bound by the terms of this Software Use Agreement (“Agreement”). Absent your agreement to the terms below, you (the “End User”) have no rights to hold or use the Software whatsoever.
+
+Harvard agrees to grant hereunder a limited non-exclusive license to End User for the use of the Software in the performance of End User’s internal, non-commercial research and academic use at End User’s academic or not-for-profit research institution (“Institution”) on the following terms and conditions:
+
+1. **NO COMMERCIAL USE.** End User shall not use the Software for Commercial use and any such use of the Software is expressly prohibited. “Commercial use” includes, but is not limited to, (i) use of the Software in fee-for-service arrangements, (ii) use of the Software by core facilities or laboratories to provide research services to (or in collaboration with) for-profit third parties for a fee, and (iii) use of the Software in industry-sponsored and/or collaborative research projects in which any commercial rights are granted to the sponsor or collaborator. If End User wishes to use the Software for Commercial use, End User must execute a separate license agreement with Harvard.
+
+_Requests for use of the Software for Commercial use, please contact:_
+
+Office of Technology Development
+Harvard University
+Smith Campus Center, Suite 727E
+1350 Massachusetts Avenue Cambridge, MA 02138 USA Telephone: (617) 495-3067
+Facsimile: (617) 495-9568
+E-mail: otd@harvard.edu
+
+2. **OWNERSHIP AND COPYRIGHT NOTICE.** Harvard owns all intellectual property in the Software. End User shall gain no ownership to the Software.  End User shall not remove or delete, and shall retain in the Software (including in any modifications to the Software and in any Derivative Works), the copyright, trademark, or other notices pertaining to Software as provided with the Software.
+
+3. **DERIVATIVE WORKS.** End User may create use and distribute Derivative Works, as such term is defined under U.S. copyright laws, provided that any such Derivative Works shall be restricted to non-commercial research and academic use. End User may not distribute Derivative Works to any for-profit third parties for Commercial use.
+
+4. **FEEDBACK.** In order to improve the Software, comments from End Users may be useful. End User agrees to provide Harvard with feedback on the End User’s use of the Software (e.g., any bugs in the Software, the user experience, etc.).  Harvard is permitted to use such information provided by End User in making changes and improvements to the Software without compensation or accounting to End User.
+
+5. **NON ASSERT.** End User acknowledges that Harvard may develop modifications to the Software that may be based on the feedback provided by End User under Section 5 above. Harvard shall not be restricted in any way by End User regarding its use of such information.  End User acknowledges the right of Harvard to prepare, publish, display, reproduce, transmit and or use modifications to the Software that may be substantially similar or functionally equivalent to End User’s modifications and/or improvements if any.  In the event that End User obtains patent protection for any modification or improvement to Software, End User agrees not to allege or enjoin infringement of End User’s patent against Harvard, or any of the researchers, medical or research staff, officers, directors and employees of those institutions.
+
+6. **PUBLICATION & ATTRIBUTION.** End User has the right to publish, present, or share results from the use of the Software.  In accordance with customary academic practice, End User will acknowledge Harvard as the provider of the Software and may cite the relevant reference(s) from the following list of publications:
+
+Wiltschko AB, Johnson MJ, Iurilli G, Peterson RE, Katon JE, Pashkovski SL, Abraira VE, Adams RP and Datta SR. (2015). Mapping sub-second Structure in Behavior. Neuron, 88:1121.
+
+Markowitz JE, Gillis WF, Beron CC, Neufeld SQ, Robertson K, Bhagat ND, Peterson RE, Peterson E, Hyun M, Linderman SW, Sabatini BL and Datta SR. (2018). The striatum organizes 3D behavior via moment-to-moment action selection. Cell, in press.
+
+7.** NO WARRANTIES.** THE SOFTWARE IS PROVIDED "AS IS." TO THE FULLEST EXTENT PERMITTED BY LAW, HARVARD HEREBY DISCLAIMS ALL WARRANTIES OF ANY KIND (EXPRESS, IMPLIED OR OTHERWISE) REGARDING THE SOFTWARE, INCLUDING BUT NOT LIMITED TO ANY IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, OWNERSHIP, AND NON-INFRINGEMENT.  HARVARD MAKES NO WARRANTY ABOUT THE ACCURACY, RELIABILITY, COMPLETENESS, TIMELINESS, SUFFICIENCY OR QUALITY OF THE SOFTWARE.  HARVARD DOES NOT WARRANT THAT THE SOFTWARE WILL OPERATE WITHOUT ERROR OR INTERRUPTION.
+
+8. **LIMITATIONS OF LIABILITY AND REMEDIES.** USE OF THE SOFTWARE IS AT END USER’S OWN RISK. IF END USER IS DISSATISFIED WITH THE SOFTWARE, ITS EXCLUSIVE REMEDY IS TO STOP USING IT.  IN NO EVENT SHALL HARVARD BE LIABLE TO END USER OR ITS INSTITUTION, IN CONTRACT, TORT OR OTHERWISE, FOR ANY DIRECT, INDIRECT, SPECIAL, INCIDENTAL, CONSEQUENTIAL, PUNITIVE OR OTHER DAMAGES OF ANY KIND WHATSOEVER ARISING OUT OF OR IN CONNECTION WITH THE SOFTWARE, EVEN IF HARVARD IS NEGLIGENT OR OTHERWISE AT FAULT, AND REGARDLESS OF WHETHER HARVARD IS ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+
+9. **INDEMNIFICATION.** To the extent permitted by law, End User shall indemnify, defend and hold harmless Harvard, their corporate affiliates, current or future directors, trustees, officers, faculty, medical and professional staff, employees, students and agents and their respective successors, heirs and assigns (the "Indemnitees"), against any liability, damage, loss or expense (including reasonable attorney's fees and expenses of litigation) incurred by or imposed upon the Indemnitees or any one of them in connection with any claims, suits, actions, demands or judgments arising from End User’s breach of this Agreement or its Institution’s use of the Software except to the extent caused by the gross negligence or willful misconduct of Harvard. This indemnification provision shall survive expiration or termination of this Agreement.
+
+10. **GOVERNING LAW.** This Agreement shall be construed and governed by the laws of the Commonwealth of Massachusetts regardless of otherwise applicable choice of law standards.
+
+11. **NON-USE OF NAME.**  Nothing in this License and Terms of Use shall be construed as granting End Users or their Institutions any rights or licenses to use any trademarks, service marks or logos associated with the Software.  End User may not use the terms “Harvard” (or a substantially similar term) in any way that is inconsistent with the permitted uses described herein. End Users may not use any name or emblem of Harvard or any of its schools or subdivisions for any purpose, or to falsely suggest any relationship between End User (or its Institution) and Harvard, or in any manner that would infringe or violate any of its rights.
+
+12. End User represents and warrants that it has the legal authority to enter into this License and Terms of Use on behalf of itself and its Institution.
+
+                                                                  ***      
+
+

moseq2_model-1.4.0/PKG-INFO

@@ -0,0 +1,84 @@
+Metadata-Version: 2.4
+Name: moseq2-model
+Version: 1.4.0
+Summary: Modeling for the best
+Author: Datta Lab
+License: LicenseRef-Harvard-Academic-Non-Commercial
+Platform: any
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Requires-Dist: jax-moseq>=0.3.5
+Requires-Dist: jax
+Requires-Dist: numpy
+Requires-Dist: scipy
+Requires-Dist: h5py>=3
+Requires-Dist: click>=8
+Requires-Dist: joblib
+Requires-Dist: cytoolz
+Requires-Dist: tqdm>=4.48.0
+Requires-Dist: matplotlib
+Requires-Dist: ruamel.yaml>=0.17
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: license
+Dynamic: license-file
+Dynamic: platform
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# MoSeq2-model 
+
+[![codecov](https://codecov.io/gh/dattalab/moseq2_model/branch/master/graph/badge.svg?token=q9xxVhps5o)](https://codecov.io/gh/dattalab/moseq2_model)
+[![DOI](https://zenodo.org/badge/81390468.svg)](https://zenodo.org/badge/latestdoi/81390468)
+
+# Installation
+
+```bash
+pip install moseq2-model
+```
+
+Python 3.10 or newer, on Linux. macOS and Windows support is being verified.
+
+Most people install the whole pipeline in one step with `pip install
+moseq2-app`, which brings this package with it. See [the installation
+guide](https://github.com/dattalab/moseq2-app/wiki/MoSeq2-Installation).
+
+# [Documentation: MoSeq2 Wiki](https://github.com/dattalab/moseq2-app/wiki)
+You can find more information about MoSeq Pipeline, installation, step-by-step instructions, documentation for Command Line Interface(CLI), tutorials etc in [MoSeq2 Wiki](https://github.com/dattalab/moseq2-app/wiki).
+
+You can run `moseq2-model --version` to check the current version and `moseq2-model --help` to see all the commands.
+```bash
+Usage: moseq2-model [OPTIONS] COMMAND [ARGS]...
+
+Options:
+  --version  Show the version and exit.  [default: False]
+  --help     Show this message and exit.  [default: False]
+
+Commands:
+  count-frames  Counts number of frames in given h5 file (pca_scores)
+  kappa-scan    Batch fit multiple models scanning over different syllable...
+  learn-model   Trains ARHMM on PCA Scores with given training parameters
+```
+
+# Community Support and Contributing
+- Please join [![MoSeq Slack Channel](https://img.shields.io/badge/slack-MoSeq-blue.svg?logo=slack)](https://moseqworkspace.slack.com) to post questions and interactive with MoSeq developers and users.
+- If you encounter bugs, errors or issues, please submit a Bug report [here](https://github.com/dattalab/moseq2-app/issues/new/choose). We encourage you to check out the [troubleshooting and tips section](https://github.com/dattalab/moseq2-app/wiki/Troubleshooting-and-Tips) and search your issues in [the existing issues](https://github.com/dattalab/moseq2-app/issues) first.   
+- If you want to see certain features in MoSeq or you have new ideas, please submit a Feature request [here](https://github.com/dattalab/moseq2-app/issues/new/choose).
+- If you want to contribute to our codebases, please check out our [Developer Guidelines](https://github.com/dattalab/moseq2-app/wiki/MoSeq-Developer-Guidelines).
+- Please tell us what you think by filling out [this user survey](https://forms.gle/FbtEN8E382y8jF3p6).
+
+# License
+MoSeq is freely available for academic use under a license provided by Harvard University. Please refer to the license file for details. If you are interested in using MoSeq for commercial purposes please contact Bob Datta directly at srdatta@hms.harvard.edu, who will put you in touch with the appropriate people in the Harvard Technology Transfer office.

moseq2_model-1.4.0/README.md

@@ -0,0 +1,43 @@
+# MoSeq2-model 
+
+[![codecov](https://codecov.io/gh/dattalab/moseq2_model/branch/master/graph/badge.svg?token=q9xxVhps5o)](https://codecov.io/gh/dattalab/moseq2_model)
+[![DOI](https://zenodo.org/badge/81390468.svg)](https://zenodo.org/badge/latestdoi/81390468)
+
+# Installation
+
+```bash
+pip install moseq2-model
+```
+
+Python 3.10 or newer, on Linux. macOS and Windows support is being verified.
+
+Most people install the whole pipeline in one step with `pip install
+moseq2-app`, which brings this package with it. See [the installation
+guide](https://github.com/dattalab/moseq2-app/wiki/MoSeq2-Installation).
+
+# [Documentation: MoSeq2 Wiki](https://github.com/dattalab/moseq2-app/wiki)
+You can find more information about MoSeq Pipeline, installation, step-by-step instructions, documentation for Command Line Interface(CLI), tutorials etc in [MoSeq2 Wiki](https://github.com/dattalab/moseq2-app/wiki).
+
+You can run `moseq2-model --version` to check the current version and `moseq2-model --help` to see all the commands.
+```bash
+Usage: moseq2-model [OPTIONS] COMMAND [ARGS]...
+
+Options:
+  --version  Show the version and exit.  [default: False]
+  --help     Show this message and exit.  [default: False]
+
+Commands:
+  count-frames  Counts number of frames in given h5 file (pca_scores)
+  kappa-scan    Batch fit multiple models scanning over different syllable...
+  learn-model   Trains ARHMM on PCA Scores with given training parameters
+```
+
+# Community Support and Contributing
+- Please join [![MoSeq Slack Channel](https://img.shields.io/badge/slack-MoSeq-blue.svg?logo=slack)](https://moseqworkspace.slack.com) to post questions and interactive with MoSeq developers and users.
+- If you encounter bugs, errors or issues, please submit a Bug report [here](https://github.com/dattalab/moseq2-app/issues/new/choose). We encourage you to check out the [troubleshooting and tips section](https://github.com/dattalab/moseq2-app/wiki/Troubleshooting-and-Tips) and search your issues in [the existing issues](https://github.com/dattalab/moseq2-app/issues) first.   
+- If you want to see certain features in MoSeq or you have new ideas, please submit a Feature request [here](https://github.com/dattalab/moseq2-app/issues/new/choose).
+- If you want to contribute to our codebases, please check out our [Developer Guidelines](https://github.com/dattalab/moseq2-app/wiki/MoSeq-Developer-Guidelines).
+- Please tell us what you think by filling out [this user survey](https://forms.gle/FbtEN8E382y8jF3p6).
+
+# License
+MoSeq is freely available for academic use under a license provided by Harvard University. Please refer to the license file for details. If you are interested in using MoSeq for commercial purposes please contact Bob Datta directly at srdatta@hms.harvard.edu, who will put you in touch with the appropriate people in the Harvard Technology Transfer office.

moseq2_model-1.4.0/moseq2_model/__init__.py

@@ -0,0 +1 @@
+__version__ = "1.4.0"

moseq2_model-1.4.0/moseq2_model/_yaml.py

@@ -0,0 +1,46 @@
+"""YAML reading and writing, in one place.
+
+``ruamel.yaml`` removed the module-level ``safe_load``/``safe_dump`` functions in
+0.18 in favour of instantiating a parser. Routing every call through here keeps
+that detail in one file rather than at each of the call sites.
+
+The parser stays ``ruamel``, deliberately. ``ruamel``'s safe loader is YAML 1.2;
+PyYAML is YAML 1.1 and reads unquoted ``yes``/``no``/``on``/``off`` as booleans
+and ``06:00:00`` as a sexagesimal integer. moseq2 config files contain both
+shapes -- ``wall_time: 06:00:00`` is interpolated straight into an ``sbatch -t``
+argument by ``util.get_parameter_strings`` -- so switching parsers would quietly
+change how existing user configs are read.
+"""
+
+import io
+
+from ruamel.yaml import YAML
+
+# ``pure=True`` selects the Python implementation, matching what the module-level
+# ``safe_load`` used and keeping behaviour identical whether or not the C
+# extension is available.
+_yaml = YAML(typ="safe", pure=True)
+
+
+def read_yaml(path_or_stream):
+    """Parse a YAML file, given a path or an open stream."""
+    if hasattr(path_or_stream, "read"):
+        return _yaml.load(path_or_stream)
+    with open(path_or_stream, "r") as handle:
+        return _yaml.load(handle)
+
+
+def write_yaml(data, path_or_stream):
+    """Write ``data`` as YAML, given a path or an open stream."""
+    if hasattr(path_or_stream, "write"):
+        _yaml.dump(data, path_or_stream)
+        return
+    with open(path_or_stream, "w") as handle:
+        _yaml.dump(data, handle)
+
+
+def dump_yaml(data):
+    """Serialize ``data`` to a YAML string."""
+    stream = io.StringIO()
+    _yaml.dump(data, stream)
+    return stream.getvalue()

moseq2_model-1.4.0/moseq2_model/cli.py

@@ -0,0 +1,326 @@
+"""
+CLI for modeling the data using AR-HMM.
+"""
+
+import os
+import click
+from os.path import join
+from moseq2_model.util import count_frames as count_frames_wrapper
+from moseq2_model.helpers.wrappers import (
+    learn_model_wrapper,
+    kappa_scan_fit_models_wrapper,
+    apply_model_wrapper,
+)
+
+orig_init = click.core.Option.__init__
+
+
+def new_init(self, *args, **kwargs):
+    orig_init(self, *args, **kwargs)
+    self.show_default = True
+
+
+click.core.Option.__init__ = new_init
+
+
+@click.group()
+@click.version_option()
+def cli():
+    pass
+
+
+@cli.command(
+    name="count-frames", help="Count the number of frames in given h5 file (pca_scores)"
+)
+@click.argument("input_file", type=click.Path(exists=True))
+@click.option(
+    "--var-name",
+    type=str,
+    default="scores",
+    help="Variable name in input file with PCs",
+)
+def count_frames(input_file, var_name):
+    # Count the number of frames in the input file.
+    count_frames_wrapper(input_file=input_file, var_name=var_name)
+
+
+def modeling_parameters(function):
+
+    function = click.option(
+        "--hold-out",
+        "-h",
+        is_flag=True,
+        default=False,
+        help="Hold out one fold (set by nfolds) for computing heldout likelihood",
+    )(function)
+    function = click.option(
+        "--hold-out-seed",
+        type=int,
+        default=-1,
+        help="Random seed for holding out data (set for reproducibility)",
+    )(function)
+    function = click.option(
+        "--nfolds", type=int, default=5, help="Number of folds for split"
+    )(function)
+    function = click.option(
+        "--ncpus",
+        "-c",
+        type=int,
+        default=0,
+        help="Number of cores to use for resampling",
+    )(function)
+    function = click.option(
+        "--num-iter",
+        "-n",
+        type=int,
+        default=100,
+        help="Number of interations to resample model",
+    )(function)
+    function = click.option(
+        "--var-name",
+        type=str,
+        default="scores",
+        help="Variable name in input file with PCs",
+    )(function)
+    function = click.option(
+        "--e-step",
+        is_flag=True,
+        default=False,
+        help="Compute the expected state sequence for each recordings",
+    )(function)
+    function = click.option(
+        "--save-every",
+        "-s",
+        type=int,
+        default=-1,
+        help="Increment to save labels and model object (-1 for just last)",
+    )(function)
+    function = click.option(
+        "--save-model",
+        type=bool,
+        default=True,
+        help="Save model object at the end of training",
+    )(function)
+    function = click.option(
+        "--max-states", "-m", type=int, default=100, help="Maximum number of states"
+    )(function)
+    function = click.option(
+        "--npcs", type=int, default=10, help="Number of PCs to use"
+    )(function)
+    function = click.option(
+        "--whiten",
+        "-w",
+        type=str,
+        default="all",
+        help="Whiten PCs: (e)each session (a)ll combined or (n)o whitening",
+    )(function)
+    function = click.option(
+        "--progressbar", "-p", type=bool, default=True, help="Show model progress"
+    )(function)
+    function = click.option(
+        "--percent-split",
+        type=int,
+        default=0,
+        help="Training-validation split percentage used when not holding out data and when this parameter > 0.",
+    )(function)
+    function = click.option(
+        "--load-groups",
+        type=bool,
+        default=True,
+        help="If groups should be loaded with the PC scores.",
+    )(function)
+    function = click.option(
+        "--gamma",
+        "-g",
+        type=float,
+        default=1e3,
+        help="Gamma; hierarchical dirichlet process hyperparameter (try not to change it).",
+    )(function)
+    function = click.option(
+        "--alpha",
+        "-a",
+        type=float,
+        default=5.7,
+        help="Alpha; hierarchical dirichlet process hyperparameter (try not to change it).",
+    )(function)
+    function = click.option(
+        "--noise-level",
+        type=float,
+        default=0,
+        help="Additive white gaussian noise to input data for regularization. Not generally used",
+    )(function)
+    function = click.option(
+        "--nlags", type=int, default=3, help="Number of lags to use"
+    )(function)
+    function = click.option(
+        "--separate-trans",
+        is_flag=True,
+        default=False,
+        help="Use separate transition matrix for each group",
+    )(function)
+    function = click.option(
+        "--robust",
+        is_flag=True,
+        default=False,
+        help="Use robust AR-HMM model. More tolerant to noise",
+    )(function)
+    function = click.option(
+        "--check-every",
+        type=int,
+        default=5,
+        help="Increment to record training and validation log-likelihoods.",
+    )(function)
+
+    return function
+
+
+@cli.command(
+    name="learn-model", help="Train ARHMM on PC Scores with given training parameters"
+)
+@click.argument("input_file", type=click.Path(exists=True))
+@click.argument(
+    "dest_file", type=click.Path(file_okay=True, writable=True, resolve_path=True)
+)
+@modeling_parameters
+@click.option(
+    "--kappa",
+    "-k",
+    type=float,
+    default=None,
+    help="Kappa; hyperparameter used to set syllable duration. Larger k = longer syllable lengths",
+)
+@click.option(
+    "--checkpoint-freq",
+    type=int,
+    default=-1,
+    help="save model checkpoint every n iterations",
+)
+@click.option(
+    "--use-checkpoint",
+    is_flag=True,
+    default=False,
+    help="indicate whether to use previously saved checkpoint",
+)
+@click.option(
+    "--index",
+    "-i",
+    type=click.Path(),
+    default="",
+    help="Path to moseq2-index.yaml for group definitions",
+)
+@click.option(
+    "--default-group",
+    type=str,
+    default="n/a",
+    help="Default group name to use for separate-trans",
+)
+@click.option(
+    "--verbose",
+    "-v",
+    is_flag=True,
+    default=False,
+    help="Print syllable log-likelihoods during training.",
+)
+def learn_model(input_file, dest_file, **config_data):
+    # Train the ARHMM using PC scores located in the INPUT_FILE, and saves the model to DEST_FILE
+
+    learn_model_wrapper(input_file, dest_file, config_data)
+
+
+@cli.command(name="apply-model", help="Apply pre-trained ARHMM to PC scores.")
+@click.argument("model_file", type=click.Path(exists=True))
+@click.argument("pc_file", type=click.Path(exists=True))
+@click.argument(
+    "dest_file", type=click.Path(file_okay=True, writable=True, resolve_path=True)
+)
+@click.option(
+    "--var-name",
+    type=str,
+    default="scores",
+    help="Variable name in input file with PCs",
+)
+@click.option(
+    "--index",
+    "-i",
+    type=click.Path(),
+    default="",
+    help="Path to moseq2-index.yaml for group definitions",
+)
+@click.option(
+    "--load-groups",
+    type=bool,
+    default=True,
+    help="If groups should be loaded with the PC scores.",
+)
+def apply_model(model_file, pc_file, dest_file, **config_data):
+    # Apply the ARHMM located in MODEL_FILE to the PC scores in PC_FILE, and saves the results to DEST_FILE
+
+    apply_model_wrapper(model_file, pc_file, dest_file, config_data)
+
+
+@cli.command(
+    name="kappa-scan",
+    help="Batch train multiple model to scan over different kappa values.",
+)
+@click.argument("input_file", type=click.Path(exists=True))
+@click.argument("output_dir", type=click.Path(exists=False))
+@click.option(
+    "--index",
+    "-i",
+    type=click.Path(),
+    default="",
+    help="Path to moseq2-index.yaml for session metadata and group information",
+)
+@click.option(
+    "--out-script",
+    type=click.Path(),
+    default="train_out.sh",
+    help="Name of bash script file to save model training commands.",
+)
+@click.option(
+    "--n-models", type=int, default=10, help="Number of models to train in kappa scan."
+)
+@click.option(
+    "--prefix",
+    type=str,
+    default="",
+    help="Batch command string to prefix model training command (slurm only).",
+)
+@click.option(
+    "--cluster-type",
+    type=click.Choice(["local", "slurm"]),
+    default="local",
+    help="Platform to train models on",
+)
+@click.option(
+    "--scan-scale",
+    type=click.Choice(["log", "linear"]),
+    default="log",
+    help="Scale to scan kappa values at.",
+)
+@click.option(
+    "--min-kappa",
+    type=float,
+    default=None,
+    help="Minimum kappa value to begin scan from.",
+)
+@click.option(
+    "--max-kappa", type=float, default=None, help="Maximum kappa value to end scan on."
+)
+@click.option("--memory", type=str, default="5GB", help="RAM (slurm only)")
+@click.option("--wall-time", type=str, default="3:00:00", help="Wall time (slurm only)")
+@click.option(
+    "--partition", type=str, default="short", help="Partition name (slurm only)"
+)
+@click.option("--get-cmd", is_flag=True, default=False, help="Print scan command strings.")
+@click.option("--run-cmd", is_flag=True, default=False, help="Run scan command strings.")
+@modeling_parameters
+def kappa_scan_fit_models(input_file, output_dir, **config_data):
+    # Scan through the kappa hyperparameter to find the kappa that best matches the changepoint duration distribution.
+
+    config_data["out_script"] = join(output_dir, config_data["out_script"])
+    kappa_scan_fit_models_wrapper(input_file, config_data, output_dir)
+
+
+if __name__ == "__main__":
+    cli()

moseq2_model-1.4.0/moseq2_model/gui.py

@@ -0,0 +1,104 @@
+"""
+GUI front-end functions for training ARHMM.
+"""
+
+from moseq2_model.cli import learn_model, kappa_scan_fit_models
+from os.path import dirname, join, exists
+from moseq2_model.helpers.wrappers import (
+    learn_model_wrapper,
+    kappa_scan_fit_models_wrapper,
+    apply_model_wrapper,
+)
+from moseq2_model._yaml import read_yaml
+
+
+def learn_model_command(progress_paths, get_cmd=True, verbose=False):
+    """
+    Train ARHMM from within a Jupyter notebook using parameters specified in the notebook.
+
+    Args:
+    progress_paths (dict): notebook progress dict that contains paths to the pc scores, config, and index files.
+    get_cmd (bool): flag to return the kappa scan learn-model commands.
+    verbose (bool): compute modeling summary - can slow down training.
+
+    Returns:
+    None or kappa scan command
+    """
+
+    # Load proper input variables
+    input_file = progress_paths["scores_path"]
+    dest_file = progress_paths["model_path"]
+    config_file = progress_paths["config_file"]
+    index = progress_paths["index_file"]
+    output_dir = progress_paths["base_model_path"]
+
+    assert exists(
+        input_file
+    ), "PCA Score not found; set the correct path in progress_paths['scores_path']"
+    assert exists(
+        config_file
+    ), "Config file not found; set the correct path in progress_paths['config_file']"
+    assert exists(
+        index
+    ), "Index file not found; set the correct path in progress_paths['index_file']"
+
+    with open(config_file, "r") as f:
+        config_data = read_yaml(f)
+
+    # Get default CLI params
+    params = {tmp.name: tmp.default for tmp in learn_model.params if not tmp.required}
+    # merge default params and config data, preferring values in config data
+    config_data = {**params, **config_data}
+
+    config_data["verbose"] = verbose
+    config_data["index"] = index
+
+    if output_dir is None:
+        print("Output directory not specified, saving models to base directory.")
+        output_dir = dirname(index)
+
+    if config_data["kappa"] == "scan":
+        assert config_data["scan_scale"] in (
+            "log",
+            "linear",
+        ), 'scan scale must be "log" or "linear"'
+
+        # Get default CLI params
+        params = {
+            tmp.name: tmp.default
+            for tmp in kappa_scan_fit_models.params
+            if not tmp.required
+        }
+        # merge default params and config data, preferring values in config data
+        config_data = {**params, **config_data}
+
+        config_data["out_script"] = join(output_dir, config_data["out_script"])
+        config_data["get_cmd"] = get_cmd
+
+        command = kappa_scan_fit_models_wrapper(input_file, config_data, output_dir)
+        return command
+    else:
+        learn_model_wrapper(input_file, dest_file, config_data)
+
+
+def apply_model_command(progress_paths, model_file):
+    """Apply a pre-trained ARHMM to a new dataset from within a Jupyter notebook.
+
+    Args:
+        progress_paths (dict): notebook progress dict that contains paths to the pc scores, config, and index files.
+        model_file (str): path to the pre-trained ARHMM.
+    """
+
+    # Load proper input variables
+    pc_file = progress_paths["scores_path"]
+    dest_file = progress_paths["model_path"]
+    config_file = progress_paths["config_file"]
+    index = progress_paths["index_file"]
+    output_dir = progress_paths["base_model_path"]
+
+    # load config data
+    with open(config_file, "r") as f:
+        config_data = read_yaml(f)
+
+    # apply model to data
+    apply_model_wrapper(model_file, pc_file, dest_file, config_data)