Reorganize codebase (#7)

.pre-commit-config.yaml

@@ -1,16 +1,21 @@
 repos:
 - repo: https://github.com/hadialqattan/pycln
-  rev: "v2.4.0"
+  rev: "v2.5.0"
   hooks:
   - id: pycln
     args:
     - --all
 
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.8.2
+  rev: v0.8.6
   hooks:
-    - id: ruff-format
     - id: ruff
+      name: ruff linting
+    - id: ruff
+      name: ruff import sorting
+      args: [ --select, I, --fix ]
+    - id: ruff-format
+      name: ruff formatting
 
 - repo: https://github.com/pre-commit/pre-commit-hooks
   rev: 'v5.0.0'
@@ -21,7 +26,7 @@ repos:
   - id: check-merge-conflict
 
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: 'v1.13.0'
+  rev: 'v1.14.1'
   hooks:
   - id: mypy
     additional_dependencies:

pyproject.toml

@@ -78,6 +78,7 @@ fallback_version = '0.0.0'
 
 [tool.ruff]
 line-length = 100
+src = ["src"]
 
 [tool.ruff.lint]
 select = [
@@ -99,7 +100,7 @@ ignore = [
 ]
 
 [tool.ruff.lint.per-file-ignores]
-"src/furax/_base/core.py" = ['E743']
+"src/furax/core/_base.py" = ['E743']
 "src/furax/landscapes.py" = ['UP007']
 
 [tool.ruff.format]

src/furax/__init__.py

@@ -1,14 +1,53 @@
-from . import tree
-from ._base.axes import MoveAxisOperator, RavelOperator, ReshapeOperator
-from ._base.config import Config
+from ._config import Config
+from .core import (
+    AbstractLinearOperator,
+    BlockColumnOperator,
+    BlockDiagonalOperator,
+    BlockRowOperator,
+    BroadcastDiagonalOperator,
+    DenseBlockDiagonalOperator,
+    DiagonalOperator,
+    HomothetyOperator,
+    IdentityOperator,
+    IndexOperator,
+    MoveAxisOperator,
+    RavelOperator,
+    ReshapeOperator,
+    SymmetricBandToeplitzOperator,
+    diagonal,
+    lower_triangular,
+    negative_semidefinite,
+    orthogonal,
+    positive_semidefinite,
+    square,
+    symmetric,
+    upper_triangular,
+)
 
 __all__ = [
-    # _base.config
-    'Config',
-    # tree
-    'tree',
-    # _base.axes
+    # core
+    'AbstractLinearOperator',
+    'IdentityOperator',
+    'HomothetyOperator',
     'MoveAxisOperator',
     'RavelOperator',
     'ReshapeOperator',
+    'BlockRowOperator',
+    'BlockDiagonalOperator',
+    'BlockColumnOperator',
+    'DenseBlockDiagonalOperator',
+    'BroadcastDiagonalOperator',
+    'DiagonalOperator',
+    'IndexOperator',
+    'SymmetricBandToeplitzOperator',
+    'diagonal',
+    'lower_triangular',
+    'upper_triangular',
+    'symmetric',
+    'positive_semidefinite',
+    'negative_semidefinite',
+    'square',
+    'orthogonal',
+    # config
+    'Config',
 ]

src/furax/_base/config.py → src/furax/_config.py

@@ -7,6 +7,8 @@
 import lineax as lx
 import yaml
 
+__all__ = ['Config']
+
 
 def default_solver_callback(solution: lx.Solution) -> None:
     num_steps = solution.stats['num_steps']

src/furax/instruments/sat.py → src/furax/_instruments/sat.py

@@ -1,12 +1,11 @@
 import numpy as np
 
-from furax.detectors import DetectorArray
-from furax.landscapes import HealpixLandscape
-from furax.operators import AbstractLinearOperator
-from furax.operators.hwp import HWPOperator
-from furax.operators.polarizers import LinearPolarizerOperator
-from furax.projections import create_projection_operator
-from furax.samplings import Sampling
+from furax import AbstractLinearOperator
+from furax.obs import HWPOperator, LinearPolarizerOperator
+from furax.obs._detectors import DetectorArray
+from furax.obs._projections import create_projection_operator
+from furax.obs._samplings import Sampling
+from furax.obs.landscapes import HealpixLandscape
 
 FOV_DEG = 35
 DIAMETER = 0.42

src/furax/core/__init__.py

@@ -0,0 +1,57 @@
+from ._axes import MoveAxisOperator, RavelOperator, ReshapeOperator
+from ._base import (
+    AbstractLazyInverseOperator,
+    AbstractLazyInverseOrthogonalOperator,
+    AbstractLinearOperator,
+    AdditionOperator,
+    CompositionOperator,
+    HomothetyOperator,
+    IdentityOperator,
+    InverseOperator,
+    TransposeOperator,
+    diagonal,
+    lower_triangular,
+    negative_semidefinite,
+    orthogonal,
+    positive_semidefinite,
+    square,
+    symmetric,
+    upper_triangular,
+)
+from ._blocks import BlockColumnOperator, BlockDiagonalOperator, BlockRowOperator
+from ._dense import DenseBlockDiagonalOperator
+from ._diagonal import BroadcastDiagonalOperator, DiagonalOperator
+from ._indices import IndexOperator
+from ._toeplitz import SymmetricBandToeplitzOperator, dense_symmetric_band_toeplitz
+
+__all__ = [
+    'AbstractLinearOperator',
+    'AbstractLazyInverseOperator',
+    'AbstractLazyInverseOrthogonalOperator',
+    'TransposeOperator',
+    'InverseOperator',
+    'AdditionOperator',
+    'CompositionOperator',
+    'IdentityOperator',
+    'HomothetyOperator',
+    'diagonal',
+    'lower_triangular',
+    'upper_triangular',
+    'symmetric',
+    'positive_semidefinite',
+    'negative_semidefinite',
+    'square',
+    'orthogonal',
+    'MoveAxisOperator',
+    'RavelOperator',
+    'ReshapeOperator',
+    'BlockRowOperator',
+    'BlockDiagonalOperator',
+    'BlockColumnOperator',
+    'DenseBlockDiagonalOperator',
+    'BroadcastDiagonalOperator',
+    'DiagonalOperator',
+    'IndexOperator',
+    'SymmetricBandToeplitzOperator',
+    'dense_symmetric_band_toeplitz',
+]

src/furax/_base/axes.py → src/furax/core/_axes.py

@@ -8,7 +8,7 @@
 from jax import numpy as jnp
 from jaxtyping import Inexact, PyTree
 
-from .core import AbstractLinearOperator, IdentityOperator, TransposeOperator
+from ._base import AbstractLinearOperator, IdentityOperator, TransposeOperator
 from .rules import AbstractBinaryRule, NoReduction
 
 __all__ = [

src/furax/_base/core.py → src/furax/core/_base.py

@@ -11,7 +11,7 @@
 from jax import Array
 from jaxtyping import Inexact, PyTree, Scalar, ScalarLike
 
-from furax._base.config import Config, ConfigState
+from furax._config import Config, ConfigState
 from furax.tree import zeros_like
 
 

src/furax/_base/blocks.py → src/furax/core/_blocks.py

@@ -8,7 +8,7 @@
 from jax import Array
 from jaxtyping import Inexact, PyTree
 
-from .core import AbstractLinearOperator, AdditionOperator, IdentityOperator
+from ._base import AbstractLinearOperator, AdditionOperator, IdentityOperator
 from .rules import AbstractBinaryRule
 
 

src/furax/_base/dense.py → src/furax/core/_dense.py

@@ -6,9 +6,10 @@
 from jax import numpy as jnp
 from jaxtyping import Inexact, PyTree
 
-from furax._base.core import AbstractLinearOperator
 from furax.tree import is_leaf
 
+from ._base import AbstractLinearOperator
+
 
 class DenseBlockDiagonalOperator(AbstractLinearOperator):
     """Operator for block diagonal dense matrix operations involving pytrees.

src/furax/_base/diagonal.py → src/furax/core/_diagonal.py

@@ -7,9 +7,10 @@
 from jax import numpy as jnp
 from jaxtyping import Inexact, PyTree
 
-from furax._base.core import AbstractLazyInverseOperator, AbstractLinearOperator, diagonal
 from furax.tree import is_leaf
 
+from ._base import AbstractLazyInverseOperator, AbstractLinearOperator, diagonal
+
 
 class BroadcastDiagonalOperator(AbstractLinearOperator):
     """Class representing a generalized diagonal operation.

src/furax/_base/indices.py → src/furax/core/_indices.py

@@ -6,10 +6,12 @@
 from jax import Array
 from jaxtyping import Bool, Inexact, Integer, PyTree
 
-from .core import AbstractLinearOperator, IdentityOperator, TransposeOperator
-from .diagonal import DiagonalOperator
+from ._base import AbstractLinearOperator, IdentityOperator, TransposeOperator
+from ._diagonal import DiagonalOperator
 from .rules import AbstractBinaryRule, NoReduction
 
+__all__ = ['IndexOperator']
+
 
 class IndexOperator(AbstractLinearOperator):
     """Class for indexing operations on pytrees.

src/furax/_base/linear.py → src/furax/core/_linear.py

@@ -3,7 +3,7 @@
 from jax import Array
 from jaxtyping import Bool, PyTree
 
-from .core import AbstractLinearOperator, TransposeOperator
+from ._base import AbstractLinearOperator, TransposeOperator
 from .rules import AbstractBinaryRule
 
 

src/furax/operators/toeplitz.py → src/furax/core/_toeplitz.py

@@ -11,10 +11,11 @@
 from jax.typing import ArrayLike
 from jaxtyping import Array, Float, Inexact, PyTree
 
-from furax.operators import AbstractLinearOperator, symmetric
+from ._base import AbstractLinearOperator, symmetric
 
 __all__ = [
     'SymmetricBandToeplitzOperator',
+    'dense_symmetric_band_toeplitz',
 ]
 
 

src/furax/_base/rules.py → src/furax/core/rules.py

@@ -5,7 +5,7 @@
 import jax.numpy as jnp
 from jaxtyping import Scalar
 
-from .core import (
+from ._base import (
     AbstractLazyInverseOperator,
     AbstractLinearOperator,
     HomothetyOperator,

src/furax/toast/obs_matrix.py → src/furax/interfaces/toast/obs_matrix.py

@@ -6,7 +6,8 @@
 from jax.experimental.sparse.csr import CSR
 from jaxtyping import Inexact, PyTree
 
-from furax.operators import AbstractLinearOperator, TransposeOperator, square
+from furax import AbstractLinearOperator, square
+from furax.core import TransposeOperator
 
 
 @square