spectracles 0.5.0__tar.gz

spectracles-0.5.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Thomas Hilder
+
+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.

spectracles-0.5.0/PKG-INFO

@@ -0,0 +1,73 @@
+Metadata-Version: 2.4
+Name: spectracles
+Version: 0.5.0
+Summary: Unified spectrospatial models: glasses for your spectra.
+Author-email: Thomas Hilder <Thomas.Hilder@monash.edu>, "Andrew R. Casey" <Andrew.Casey@monash.edu>
+License-File: LICENSE
+Requires-Python: >=3.13
+Requires-Dist: equinox>=0.13.0
+Requires-Dist: jax-finufft
+Requires-Dist: jax[cpu]>=0.7.0
+Requires-Dist: matplotlib>=3.10.5
+Requires-Dist: networkx>=3.5
+Requires-Dist: tqdm>=4.67.1
+Description-Content-Type: text/markdown
+
+<div id="top"></div>
+
+<!-- PROJECT SHIELDS -->
+<!-- [![PyPI Package][pypi-shield]][pypi-url] -
+<!-- [![JOSS][JOSS-shield]][JOSS-url] -->
+
+<!-- ![GitHub Workflow Status](https://img.shields.io/github/actions/workflow/status/TomHilder/wakeflow/Tests.yml?label=tests&style=flat-square)
+[![Docs][docs-status-shield]][docs-status-url] -->
+
+<!-- PROJECT LOGO -->
+<br />
+<div align="center">
+  <a href="https://github.com/TomHilder/spectracles">
+    <img src="https://raw.githubusercontent.com/TomHilder/spectracles/main/logo.png" alt="spectrackles" width="420">
+  </a>
+
+<!--  <h3 align="center">Wakeflow</h3> -->
+
+  <p align="center">
+    Unified spectrospatial models for integral field spectroscopy in jax
+  </p>
+</div>
+
+<!-- <div align="center">
+<img src="https://raw.githubusercontent.com/TomHilder/spectracles/main/logo.png" alt="spectracles" width="420"></img>
+</div> -->
+
+## Glasses for your spectra
+
+Spectracles is a Python library for inferring properties of IFU/IFS spectra as continuous functions of sky position.
+
+It can also be used as a general-purpose statistical model library that extends [`equinox`](https://github.com/patrick-kidger/equinox) to allow for composable models that may have *coupled* parameters. It also implements some other nice features that are a bit awkward in `equinox` out of the box, like easily updating model parameters between fixed and varying.
+
+## Installation
+
+TODO
+
+## Usage
+
+TODO
+
+## Citation
+
+TODO
+
+## Help
+
+TODO
+
+### TODO
+
+- [x] Instead of replacing shared leaves with `0`, replace with some class/object instead
+- [ ] Nicer `__repr__` for `ShareModule` that actually says the memory address
+- [ ] Add memory address to the top of `print_model_tree`
+- [ ] Support tuples, lists and dicts of models as attributes of models
+- [ ] Handle non-odd number of modes
+- [ ] Write better tests
+- [ ] Rigorously type check the tests

spectracles-0.5.0/README.md

@@ -0,0 +1,58 @@
+<div id="top"></div>
+
+<!-- PROJECT SHIELDS -->
+<!-- [![PyPI Package][pypi-shield]][pypi-url] -
+<!-- [![JOSS][JOSS-shield]][JOSS-url] -->
+
+<!-- ![GitHub Workflow Status](https://img.shields.io/github/actions/workflow/status/TomHilder/wakeflow/Tests.yml?label=tests&style=flat-square)
+[![Docs][docs-status-shield]][docs-status-url] -->
+
+<!-- PROJECT LOGO -->
+<br />
+<div align="center">
+  <a href="https://github.com/TomHilder/spectracles">
+    <img src="https://raw.githubusercontent.com/TomHilder/spectracles/main/logo.png" alt="spectrackles" width="420">
+  </a>
+
+<!--  <h3 align="center">Wakeflow</h3> -->
+
+  <p align="center">
+    Unified spectrospatial models for integral field spectroscopy in jax
+  </p>
+</div>
+
+<!-- <div align="center">
+<img src="https://raw.githubusercontent.com/TomHilder/spectracles/main/logo.png" alt="spectracles" width="420"></img>
+</div> -->
+
+## Glasses for your spectra
+
+Spectracles is a Python library for inferring properties of IFU/IFS spectra as continuous functions of sky position.
+
+It can also be used as a general-purpose statistical model library that extends [`equinox`](https://github.com/patrick-kidger/equinox) to allow for composable models that may have *coupled* parameters. It also implements some other nice features that are a bit awkward in `equinox` out of the box, like easily updating model parameters between fixed and varying.
+
+## Installation
+
+TODO
+
+## Usage
+
+TODO
+
+## Citation
+
+TODO
+
+## Help
+
+TODO
+
+### TODO
+
+- [x] Instead of replacing shared leaves with `0`, replace with some class/object instead
+- [ ] Nicer `__repr__` for `ShareModule` that actually says the memory address
+- [ ] Add memory address to the top of `print_model_tree`
+- [ ] Support tuples, lists and dicts of models as attributes of models
+- [ ] Handle non-odd number of modes
+- [ ] Write better tests
+- [ ] Rigorously type check the tests

spectracles-0.5.0/pyproject.toml

@@ -0,0 +1,47 @@
+[project]
+name = "spectracles"
+version = "0.5.0"
+description = "Unified spectrospatial models: glasses for your spectra."
+readme = "README.md"
+authors = [
+    { name = "Thomas Hilder", email = "Thomas.Hilder@monash.edu" },
+    { name = "Andrew R. Casey", email = "Andrew.Casey@monash.edu" },
+]
+requires-python = ">=3.13"
+dependencies = [
+    "equinox>=0.13.0",
+    "jax-finufft",
+    "jax[cpu]>=0.7.0",
+    "matplotlib>=3.10.5",
+    "networkx>=3.5",
+    "tqdm>=4.67.1",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build]
+include = ["src/spectracles/py.typed"]
+exclude = ["**/__pycache__/**"]
+
+[tool.hatch.build.targets.sdist]
+include = ["src/**", "README.md", "LICENSE"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/spectracles"]
+
+[tool.ruff]
+line-length = 99 # Sorry PEP8 I'm a rebel
+
+[tool.uv.sources]
+jax-finufft = { git = "https://github.com/flatironinstitute/jax-finufft.git" }
+
+[dependency-groups]
+dev = [
+    "mypy>=1.17.1",
+    "pytest>=8.4.1",
+    "pytest-cov>=6.2.1",
+    "ruff>=0.12.8",
+    "types-tqdm>=4.67.0.20250809",
+]

spectracles-0.5.0/src/spectracles/__init__.py

@@ -0,0 +1,35 @@
+from spectracles.model.data import SpatialDataGeneric, SpatialDataLVM
+from spectracles.model.io import load_model, save_model
+from spectracles.model.kernels import Kernel, Matern12, Matern32, Matern52, SquaredExponential
+from spectracles.model.parameter import AnyParameter, ConstrainedParameter, Parameter, l_bounded
+from spectracles.model.share_module import build_model
+from spectracles.model.spatial import FourierGP, PerSpaxel, SpatialModel
+from spectracles.model.spectral import Constant, Gaussian, SpectralSpatialModel
+from spectracles.optimise.opt_frame import OptimiserFrame
+from spectracles.optimise.opt_schedule import OptimiserSchedule, PhaseConfig
+
+__all__ = [
+    "FourierGP",
+    "Gaussian",
+    "SpatialDataGeneric",
+    "SpatialDataLVM",
+    "PerSpaxel",
+    "Constant",
+    "Kernel",
+    "Matern12",
+    "Matern32",
+    "Matern52",
+    "SquaredExponential",
+    "build_model",
+    "SpatialModel",
+    "SpectralSpatialModel",
+    "Parameter",
+    "ConstrainedParameter",
+    "AnyParameter",
+    "l_bounded",
+    "OptimiserFrame",
+    "save_model",
+    "load_model",
+    "PhaseConfig",
+    "OptimiserSchedule",
+]

spectracles-0.5.0/src/spectracles/model/__init__.py

@@ -0,0 +1 @@
+"""model - subpackage for modelling framework built on jax and equinox."""

spectracles-0.5.0/src/spectracles/model/data.py

@@ -0,0 +1,27 @@
+"""data.py - Data structures used as arguments for model evaluations/predictions."""
+
+import jax.numpy as jnp
+from equinox import Module, field
+from jaxtyping import Array
+
+
+def convert_to_flat_array(array: Array) -> Array:
+    return jnp.asarray(array).flatten()
+
+
+class SpatialDataGeneric(Module):
+    x: Array = field(converter=convert_to_flat_array)
+    y: Array = field(converter=convert_to_flat_array)
+    idx: Array = field(converter=convert_to_flat_array)
+
+
+class SpatialDataLVM(Module):
+    x: Array = field(converter=convert_to_flat_array)
+    y: Array = field(converter=convert_to_flat_array)
+    idx: Array = field(converter=convert_to_flat_array)
+    tile_idx: Array = field(converter=convert_to_flat_array)
+    ifu_idx: Array = field(converter=convert_to_flat_array)
+
+
+# TODO: This is bad. We want users to be able to write their own SpatialDataFoo class, but with the current setup the typing doesn't strictly allow this. Really this should be a Protocol or subclassing sitation. Not worth refactoring right now.
+SpatialData = SpatialDataGeneric | SpatialDataLVM

spectracles-0.5.0/src/spectracles/model/graph.py

@@ -0,0 +1,84 @@
+"""graph.py - Graph utilities for the modelling framework."""
+
+from contextlib import contextmanager
+
+import matplotlib as mpl
+from networkx import DiGraph
+
+DEFAULT_NX_KWDS = {
+    "node_size": 8_000,
+    "node_color": "white",
+    "edgecolors": "black",
+    "linewidths": 2,
+    "arrowsize": 20,
+    "font_size": 8,
+    "font_color": "black",
+}
+
+
+@contextmanager
+def temporarily_disable_tex():
+    prev_setting = mpl.rcParams["text.usetex"]
+    mpl.rcParams["text.usetex"] = False
+    try:
+        yield
+    finally:
+        mpl.rcParams["text.usetex"] = prev_setting
+
+
+def print_graph(graph: DiGraph, root_id: int, indent: str = "", is_last: bool = True) -> None:
+    # Get info for current node
+    node_data = graph.nodes[root_id]
+    name = node_data["name"]
+    node_type = node_data["type"]
+    # Format display text
+    if name is None:
+        # Root module without parent attribute name
+        display_text = node_type
+    else:
+        # Regular format: "name (Type)"
+        display_text = f"{name} ({node_type})"
+    # Print this node
+    print(f"{indent}{'└── ' if is_last else '├── '}{display_text}")
+    # Find child nodes (those pointing to this node)
+    children = []
+    for src, dst in graph.edges():
+        if src == root_id:
+            children.append(dst)
+    # Recurse for each child
+    new_indent = indent + ("    " if is_last else "│   ")
+    for i, child_id in enumerate(children):
+        is_last_child = i == len(children) - 1
+        print_graph(graph, child_id, new_indent, is_last_child)
+
+
+def layered_hierarchy_pos(G, root, total_width=1.0, vert_gap=0.2):
+    """NOTE: AI generated function"""
+    from collections import defaultdict, deque
+
+    levels = defaultdict(list)
+    visited = set()
+    queue = deque([(root, 0)])
+    max_level = 0
+
+    while queue:
+        node, level = queue.popleft()
+        if node in visited:
+            continue
+        visited.add(node)
+        levels[level].append(node)
+        max_level = max(max_level, level)
+        for child in G.successors(node):
+            queue.append((child, level + 1))
+
+    pos = {}
+    for level in range(max_level + 1):
+        nodes = levels[level]
+        n = len(nodes)
+        gap = total_width / (n + 1)
+        for i, node in enumerate(nodes):
+            x = (i + 1) * gap
+            y = -level * vert_gap
+            pos[node] = (x, y)
+
+    return pos

spectracles-0.5.0/src/spectracles/model/io.py

@@ -0,0 +1,78 @@
+"""io.py - Model object serialisation and deserialisation."""
+
+from pathlib import Path
+from warnings import catch_warnings, filterwarnings
+
+from dill import PicklingWarning, dump, load  # type: ignore[import]
+
+from spectracles.model.share_module import ShareModule
+
+MODELFILE_EXT = ".model"
+
+
+def save_model(model: ShareModule, file: Path, overwrite: bool = False, **dump_kwargs):
+    """
+    Save a model to a file.
+
+    Args:
+        model (ShareModule): The model to save.
+        file (Path): The file to save the model to.
+        overwrite (bool): Whether to overwrite the file if it exists. Defaults to False.
+        dump_kwargs: Additional arguments to pass to dill.dump.
+
+    Raises:
+        FileExistsError: If the file already exists and overwrite is False.
+        TypeError: If the model is not of type ShareModule.
+        Exception: If there is a warning when saving the model.
+    """
+    # Ensure always the same file extension
+    file = file.with_suffix(MODELFILE_EXT)
+    # Check that the file doesn't exist already
+    if file.exists():
+        if not overwrite:
+            raise FileExistsError("File already exists. Overwrite with overwrite=True.")
+    # Check that model has the right type
+    if not isinstance(model, ShareModule):
+        raise TypeError(
+            "model must be type ShareModule. Saving sub-components of a model is not currenly supported. Saving a model not instantiated via build_model will never be supported."
+        )
+    # We want dill warnings to be exceptions
+    with catch_warnings():
+        filterwarnings("error")
+        # Open file in write and bytes mode
+        with open(file, "wb") as f:
+            try:
+                dump(model, f, **dump_kwargs)
+            except PicklingWarning:
+                raise Exception(
+                    "Above warning raised by dill when saving. Likely, you need to move your model class into it's own file and import it."
+                )
+
+
+def load_model(file: Path, **load_kwargs) -> ShareModule:
+    """
+    Load a model from a file.
+
+    Args:
+        file (Path): The file to load the model from.
+        load_kwargs: Additional arguments to pass to dill.load.
+
+    Returns:
+        ShareModule: The loaded model.
+
+    Raises:
+        FileNotFoundError: If the file does not exist.
+        TypeError: If the loaded model is not of type ShareModule.
+    """
+    # Ensure always the same file extension
+    file = file.with_suffix(MODELFILE_EXT)
+    # Check that the file exists
+    if not file.exists():
+        raise FileNotFoundError(f"File {file} does not exist.")
+    # Open file in read and bytes mode
+    with open(file, "rb") as f:
+        model = load(f, **load_kwargs)
+    # Check that the model is of the right type
+    if not isinstance(model, ShareModule):
+        raise TypeError("Loaded from file successfully, but model is not of type ShareModule.")
+    return model

spectracles-0.5.0/src/spectracles/model/kernels.py

@@ -0,0 +1,90 @@
+"""kernels.py - Kernel classes that implement various covariance functions expressed by their power spectral density or 'feature weights' for use with Fourier-accelerated GP models."""
+
+from abc import abstractmethod
+
+import jax.numpy as jnp
+from equinox import Module
+from jaxtyping import Array, ArrayLike
+
+from spectracles.model.parameter import Parameter
+
+
+def normalise_fw(fw: ArrayLike) -> Array:
+    """
+    Normalise the feature weights. Includes a factor of sqrt(2) which accounts for the
+    halving in total power incurred by enforcing the Fourier coefficients to be real via
+    conjugate symmetry.
+    """
+    power = jnp.sum(jnp.abs(fw) ** 2)
+    return jnp.sqrt(2) * fw / jnp.sqrt(power)
+
+
+def matern_kernel_fw_nd(
+    freqs: ArrayLike,
+    length: ArrayLike,
+    var: ArrayLike,
+    nu: float,
+    n: int,
+) -> Array:
+    """Square root of the PSD of the Matern kernel in n dimensions."""
+    fw = (1 + (freqs * length) ** 2) ** (-0.5 * (nu + n / 2))
+    return jnp.sqrt(var) * normalise_fw(fw)
+
+
+class Kernel(Module):
+    # All kernels should have a length scale and variance
+    length_scale: Parameter
+    variance: Parameter
+
+    @abstractmethod
+    def feature_weights(self, freqs: Array) -> Array:
+        pass
+
+
+class Matern12(Kernel):
+    length_scale: Parameter
+    variance: Parameter
+
+    def __init__(self, length_scale: Parameter, variance: Parameter):
+        self.length_scale = length_scale
+        self.variance = variance
+
+    def feature_weights(self, freqs: Array) -> Array:
+        return matern_kernel_fw_nd(freqs, self.length_scale.val, self.variance.val, nu=0.5, n=2)
+
+
+class Matern32(Kernel):
+    length_scale: Parameter
+    variance: Parameter
+
+    def __init__(self, length_scale: Parameter, variance: Parameter):
+        self.length_scale = length_scale
+        self.variance = variance
+
+    def feature_weights(self, freqs: Array) -> Array:
+        return matern_kernel_fw_nd(freqs, self.length_scale.val, self.variance.val, nu=1.5, n=2)
+
+
+class Matern52(Kernel):
+    length_scale: Parameter
+    variance: Parameter
+
+    def __init__(self, length_scale: Parameter, variance: Parameter):
+        self.length_scale = length_scale
+        self.variance = variance
+
+    def feature_weights(self, freqs: Array) -> Array:
+        return matern_kernel_fw_nd(freqs, self.length_scale.val, self.variance.val, nu=2.5, n=2)
+
+
+class SquaredExponential(Kernel):
+    length_scale: Parameter
+    variance: Parameter
+
+    def __init__(self, length_scale: Parameter, variance: Parameter):
+        self.length_scale = length_scale
+        self.variance = variance
+
+    def feature_weights(self, freqs: Array) -> Array:
+        fw = jnp.exp(-0.25 * freqs**2 * self.length_scale.val**2 + 1e-4)
+        return jnp.sqrt(self.variance.val) * normalise_fw(fw)