Metadata-Version: 2.4
Name: phasecurvefit
Version: 0.3.1
Summary: Walk through phase-space observations
Project-URL: Bug Tracker, https://github.com/GalacticDynamics/phasecurvefit/issues
Project-URL: Changelog, https://github.com/GalacticDynamics/phasecurvefit/releases
Project-URL: Homepage, https://github.com/GalacticDynamics/phasecurvefit
Author-email: GalacticDynamics <nstarman@users.noreply.github.com>, Nathaniel Starkman <nstarman@users.noreply.github.com>
License: MIT License
        
        Copyright (c) 2025, Nathaniel Starkman.
        All rights reserved.
        
        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:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        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.
License-File: LICENSE
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering
Classifier: Typing :: Typed
Requires-Python: >=3.12
Requires-Dist: equinox>=0.13.5
Requires-Dist: jax-tqdm>=0.4.0
Requires-Dist: jax>=0.7.2
Requires-Dist: jaxmore>=0.4.0
Requires-Dist: jaxtyping>=0.3.5
Requires-Dist: optax>=0.2.4
Requires-Dist: optional-dependencies>=0.4.0
Requires-Dist: plum-dispatch>=2.6.1
Requires-Dist: quax>=0.2.1
Requires-Dist: scipy>=1.11
Requires-Dist: zeroth>=1.0.1
Provides-Extra: all
Requires-Dist: jaxkd>=0.0.5; extra == 'all'
Requires-Dist: unxt>=1.10.4; extra == 'all'
Provides-Extra: interop
Requires-Dist: unxt>=1.10.4; extra == 'interop'
Provides-Extra: kdtree
Requires-Dist: jaxkd>=0.0.5; extra == 'kdtree'
Description-Content-Type: text/markdown

# phasecurvefit: Construct Paths through Phase-Space Points

[![PyPI version](https://img.shields.io/pypi/v/phasecurvefit.svg)](https://pypi.org/project/phasecurvefit/)
[![Python versions](https://img.shields.io/pypi/pyversions/phasecurvefit.svg)](https://pypi.org/project/phasecurvefit/)

Construct paths through phase-Space points, supporting many different
algorithms.

## Features

- **JAX-powered**: Fully compatible with JAX transformations (`jit`, `vmap`,
  `grad`)
- **GPU-ready**: Runs on CPU, GPU, or TPU via JAX
- **Type-safe**: Comprehensive (optionally runtime checked) type hints with
  `jaxtyping`
- **Pluggable metrics**: Customizable distance metrics for different physical
  interpretations
- **Pluggable query strategies**: Flexible neighbor search strategies (e.g.,
  brute-force, KD-tree) to optimize performance
- **Pluggable orderers**: One interface over multiple ordering algorithms — the
  velocity-following walk and an MST backbone for near-closed loops
- **Highly customizable ML setup and training**: Well-chosen defaults with
  highly flexible customization for specific use-cases.
- **Physical units**: Optional support via `unxt` for unit-aware calculations

## Installation

Install the core package:

```bash
pip install phasecurvefit[all]
```

Or with uv:

```bash
uv add phasecurvefit[all]
```

<details>
  <summary>from source, using uv</summary>

```bash
uv add git+https://github.com/GalacticDynamics/phasecurvefit.git@main
```

You can customize the branch by replacing `main` with any other branch name.

</details>
<details>
  <summary>building from source</summary>

```bash
cd /path/to/parent
git clone https://github.com/GalacticDynamics/phasecurvefit.git
cd phasecurvefit
uv pip install -e .  # editable mode
```

</details>

### Optional Dependencies

phasecurvefit has optional dependencies for extended functionality:

- **unxt**: Physical units support for phase-space calculations
- **tree (jaxkd)**: Spatial KD-tree queries for large datasets

Install with optional dependencies:

```bash
# pip install phasecurvefit[all]  # Install with all extras
pip install phasecurvefit[interop]  # Install with unxt for unit support
pip install phasecurvefit[kdtree]  # Install with jaxkd for KD-tree strategy
```

Or with uv:

```bash
# uv add phasecurvefit --extra all  # installs all extras
uv add phasecurvefit --extra interop
uv add phasecurvefit --extra kdtree
```

## Quick Start

```python
import jax
import jax.numpy as jnp
import phasecurvefit as pcf

# Create phase-space observations as dictionaries (Cartesian coordinates)
pos = {
    "x": jnp.array([0.0, 1.0, 2.0, 3.0, 4.0]),
    "y": jnp.array([0.0, 0.5, 1.0, 1.5, 2.0]),
}
vel = {
    "x": jnp.array([1.0, 1.0, 1.0, 1.0, 1.0]),
    "y": jnp.array([0.5, 0.5, 0.5, 0.5, 0.5]),
}

# Step 1: Order the observations (use KD-tree for spatial neighbor prefiltering)
config = pcf.WalkConfig(strategy=pcf.strats.KDTree(k=3))  # k=3 for this small dataset
result = pcf.order(pos, vel, pcf.orderers.LocalFlowOrderer(config=config))
print(result.indices)  # Initial ordering

# Step 2: Create normalizer and autoencoder
key = jax.random.key(0)
normalizer = pcf.nn.StandardScalerNormalizer(pos, vel)
autoencoder = pcf.nn.PathAutoencoder.make(
    normalizer, gamma_range=result.gamma_range, key=key
)

# Step 3: Configure and run training
train_config = pcf.nn.TrainingConfig(
    n_epochs_encoder=100,  # Encoder-only epochs
    n_epochs_both=50,  # Joint training epochs
    show_pbar=False,  # Disable progress bar
)

# Train the autoencoder
result, _, losses = pcf.nn.train_autoencoder(
    autoencoder, result, config=train_config, key=key
)

print(result.indices)  # Post-training ordering
```

### With Physical Units

When `unxt` is installed, you can use physical units throughout the workflow:

```python
import jax
import jax.numpy as jnp
import phasecurvefit as pcf
import unxt as u

# Create phase-space observations with units
pos = {
    "x": u.Q([0.0, 1.0, 2.0, 3.0, 4.0], "kpc"),
    "y": u.Q([0.0, 0.5, 1.0, 1.5, 2.0], "kpc"),
}
vel = {
    "x": u.Q([1.0, 1.0, 1.0, 1.0, 1.0], "km/s"),
    "y": u.Q([0.5, 0.5, 0.5, 0.5, 0.5], "km/s"),
}

# Step 1: Order with units (units are preserved throughout)
metric_scale = u.Q(1.0, "kpc")
config = pcf.WalkConfig(strategy=pcf.strats.KDTree(k=3))
result = pcf.order(
    pos,
    vel,
    pcf.orderers.LocalFlowOrderer(config=config, metric_scale=metric_scale),
    metadata=pcf.StateMetadata(usys=u.unitsystems.galactic),
)

# Step 2: Create normalizer and autoencoder (handles units automatically)
key = jax.random.key(0)
normalizer = pcf.nn.StandardScalerNormalizer(pos, vel)
autoencoder = pcf.nn.PathAutoencoder.make(
    normalizer, gamma_range=result.gamma_range, key=key
)
result, _, losses = pcf.nn.train_autoencoder(
    autoencoder, result, config=train_config, key=key
)
```

## Orderers

The ordering step is pluggable. Every orderer implements the same interface —
`order(positions, velocities)` — and returns an `OrderingResult` that feeds the
autoencoder unchanged, so orderers are interchangeable:

- **`LocalFlowOrderer`** — the velocity-following walk (wraps
  `walk_local_flow`). Follows a coherent flow from a start point.
- **`MSTOrderer`** — a minimum-spanning-tree backbone. It needs no start point
  (the graph diameter finds the two tips itself), which makes it ideal for
  **near-closed loops** where the velocity field reverses and a single walk
  covers only one arm. Requires the `mst` extra.

```python
import jax.numpy as jnp
import phasecurvefit as pcf

# Points along a curve
t = jnp.linspace(0.0, 1.0, 60)
pos = {"x": 10.0 * t, "y": jnp.sin(3.0 * t)}
vel = {"x": jnp.ones(60), "y": 3.0 * jnp.cos(3.0 * t)}

# The velocity-following walk, via the orderer interface
walk_orderer = pcf.orderers.LocalFlowOrderer(metric_scale=1.0, start_idx=0)
walk_result = walk_orderer.order(pos, vel)

# The MST backbone (no start point needed)
mst_orderer = pcf.orderers.MSTOrderer(k=8, jump_cap=2.0)
mst_result = pcf.order(pos, vel, mst_orderer)  # or mst_orderer.order(pos, vel)

# Either result feeds the autoencoder unchanged
print(mst_result.gamma_range)  # (-1.0, 1.0)
```

`MSTOrderer` also has opt-in velocity mechanisms (`velocity_weight`,
`sever_cos_threshold`, `orient_by_velocity`) for self-overlapping streams. See
the
[Orderers Guide](https://phasecurvefit.readthedocs.io/en/latest/guides/orderers.html)
and the
[Migration Guide](https://phasecurvefit.readthedocs.io/en/latest/migration.html).

## Distance Metrics

The algorithm supports pluggable distance metrics to control how points are
ordered. The default metric is `AlignedMomentumDistanceMetric`, which combines
spatial proximity with velocity alignment:

```python
import jax.numpy as jnp
import phasecurvefit as pcf

# Define simple Cartesian arrays (not quantities)
pos = {"x": jnp.array([0.0, 1.0, 2.0]), "y": jnp.array([0.0, 0.5, 1.0])}
vel = {"x": jnp.array([1.0, 1.0, 1.0]), "y": jnp.array([0.5, 0.5, 0.5])}

# Use default metric (AlignedMomentumDistanceMetric)
config = pcf.WalkConfig()
result = pcf.order(pos, vel, pcf.orderers.LocalFlowOrderer(config=config))
```

### Using Different Metrics

`phasecurvefit` provides three built-in metrics:

1. **AlignedMomentumDistanceMetric** (default): Combines spatial distance with
   velocity alignment (momentum-weighted nearest neighbor)
2. **FullPhaseSpaceDistanceMetric**: True 6D Euclidean distance in phase space
3. **SpatialDistanceMetric**: Pure spatial distance, ignoring velocity

```python
import jax.numpy as jnp
import phasecurvefit as pcf

# Define simple Cartesian arrays (not quantities)
pos = {"x": jnp.array([0.0, 1.0, 2.0]), "y": jnp.array([0.0, 0.5, 1.0])}
vel = {"x": jnp.array([1.0, 1.0, 1.0]), "y": jnp.array([0.5, 0.5, 0.5])}

# Pure spatial ordering (ignores velocity)
config_spatial = pcf.WalkConfig(metric=pcf.metrics.SpatialDistanceMetric())
result = pcf.order(
    pos, vel, pcf.orderers.LocalFlowOrderer(config=config_spatial, metric_scale=0.0)
)

# Full 6D phase-space distance
config_phase = pcf.WalkConfig(metric=pcf.metrics.FullPhaseSpaceDistanceMetric())
result = pcf.order(pos, vel, pcf.orderers.LocalFlowOrderer(config=config_phase))
```

### Custom Metrics

You can define custom metrics by subclassing `AbstractDistanceMetric`:

```python
import jax
import jax.numpy as jnp
import phasecurvefit as pcf


class WeightedPhaseSpaceMetric(pcf.metrics.AbstractDistanceMetric):
    """Custom weighted phase-space metric."""

    def __call__(self, current_pos, current_vel, positions, velocities, metric_scale):
        # Compute position distance
        pos_diff = jax.tree.map(jnp.subtract, positions, current_pos)
        pos_dist_sq = sum(jax.tree.leaves(jax.tree.map(jnp.square, pos_diff)))

        # Compute velocity distance
        vel_diff = jax.tree.map(jnp.subtract, velocities, current_vel)
        vel_dist_sq = sum(jax.tree.leaves(jax.tree.map(jnp.square, vel_diff)))

        # Custom weighting scheme
        return jnp.sqrt(pos_dist_sq + (metric_scale**2) * vel_dist_sq)


# Use custom metric via WalkConfig
config = pcf.WalkConfig(metric=WeightedPhaseSpaceMetric())
result = pcf.order(pos, vel, pcf.orderers.LocalFlowOrderer(config=config))
```

See the
[Metrics Guide](https://phasecurvefit.readthedocs.io/en/latest/guides/metrics.html)
for more details and examples.

## Query Strategies

The algorithm supports pluggable query strategies to control how neighbors are
found. A strategy determines which points are considered as potential next steps
in the walk.

`phasecurvefit` provides two built-in strategies:

1. **BruteForce** (default): Compute distances to all remaining points and
   select the nearest one. Efficient for small to medium datasets.
2. **KDTree**: Use spatial KD-tree prefiltering to accelerate neighbor searches
   for large datasets (requires optional `jaxkd` dependency).

### Using Built-in Strategies

```python
import jax.numpy as jnp
import phasecurvefit as pcf

# Define simple Cartesian arrays
pos = {"x": jnp.array([0.0, 1.0, 2.0]), "y": jnp.array([0.0, 0.5, 1.0])}
vel = {"x": jnp.array([1.0, 1.0, 1.0]), "y": jnp.array([0.5, 0.5, 0.5])}

# Default strategy (brute-force — no configuration needed)
config_brute = pcf.WalkConfig(strategy=pcf.strats.BruteForce())
result = pcf.order(pos, vel, pcf.orderers.LocalFlowOrderer(config=config_brute))

# KD-tree strategy for faster neighbor queries (large datasets)
config_kdtree = pcf.WalkConfig(strategy=pcf.strats.KDTree(k=2))
result = pcf.order(pos, vel, pcf.orderers.LocalFlowOrderer(config=config_kdtree))
```

### Custom Query Strategies

You can define custom strategies by subclassing `AbstractQueryStrategy`:

```python
import jax.numpy as jnp
import phasecurvefit as pcf


class SmallestIndexStrategy(pcf.strats.AbstractQueryStrategy):
    """Custom strategy: select the smallest unvisited index.

    This is a toy example showing how to implement a custom strategy.
    By returning uniform distances, argmin selects the smallest index
    deterministically. In practice, distance-based strategies like BruteForce
    are more useful.
    """

    def init(self, positions, /, *, metadata):
        """No persistent state needed."""
        return None

    def query(
        self,
        state,
        /,
        current_pos,
        current_vel,
        positions,
        velocities,
        metric_fn,
        metric_scale,
    ):
        """Return uniform distances to all points.

        Since all distances are equal, the walk algorithm's argmin will
        deterministically select the smallest unvisited index.
        """
        # Get number of points
        n_points = len(next(iter(positions.values())))

        # Return uniform distances to all points
        # argmin will pick the smallest unvisited index
        distances = jnp.ones(n_points)

        return pcf.strats.QueryResult(distances=distances, indices=None)


# Use custom strategy via WalkConfig
config = pcf.WalkConfig(strategy=SmallestIndexStrategy())
result = pcf.order(pos, vel, pcf.orderers.LocalFlowOrderer(config=config))
```

## AI Usage Disclosure

Portions of this codebase (including tests and documentation) were refactored
and generated with the assistance of Language Models. All AI contributions have
been and will continue to be reviewed and verified by the human maintainers.
