Installation
============

Requirements
------------

* Python ≥ 3.11
* `mamba <https://mamba.readthedocs.io>`_

Create and activate the environment
------------------------------------

.. code-block:: bash

   mamba env create -f environment.yml
   mamba activate sum_stat

Install the package in editable mode
-------------------------------------

.. code-block:: bash

   pip install -e ".[dev,docs]"

Verify the installation
-----------------------

.. code-block:: python

   import sum_stat
   print(sum_stat.__version__)

   import jax
   print(jax.devices())  # should show CPU or GPU

Build the documentation
-----------------------

.. code-block:: bash

   cd docs && make html
   # Open docs/_build/html/index.html

Run the tests
-------------

.. code-block:: bash

   pytest tests/ -v

   # Benchmarks only (prints timing per function):
   pytest tests/test_benchmarks.py -v -s

GPU / JAX configuration
-----------------------

By default JAX uses CPU.  To use a GPU, install the GPU version of jaxlib
in the environment:

.. code-block:: bash

   pip install --upgrade "jax[cuda12]"   # or cuda11

The package enables ``jax_enable_x64`` automatically on import so that all
floating-point operations use 64-bit precision.

Remote / HPC cluster setup
---------------------------

Use ``conda-pack`` to ship the solved environment to a cluster without
running the solver remotely.

**Step 1 — pack the environment locally:**

.. code-block:: bash

   mamba activate sum_stat
   conda-pack -o sum_stat_env.tar.gz

**Step 2 — transfer and unpack on the remote server:**

.. code-block:: bash

   scp sum_stat_env.tar.gz remote-server:$SCRATCH/
   ssh remote-server
   mkdir -p $SCRATCH/sum_stat_env
   tar -xzf $SCRATCH/sum_stat_env.tar.gz -C $SCRATCH/sum_stat_env
   source $SCRATCH/sum_stat_env/bin/activate

**Step 3 — clone the repository and install in editable mode:**

.. code-block:: bash

   git clone https://github.com/JohanComparat/sum_stat.git $SCRATCH/sum_stat
   pip install -e $SCRATCH/sum_stat --no-build-isolation

**Step 4 — verify:**

.. code-block:: bash

   python -c "import sum_stat; print(sum_stat.__version__)"
   python -c "import jax; print(jax.devices())"

See :doc:`running` for the full script reference.