Make individual files private

exponax/__init__.py

@@ -1,8 +1,10 @@
-from . import ic, metrics, nonlin_fun, normalized, poisson, stepper
-from .forced_stepper import ForcedStepper
-from .repeated_stepper import RepeatedStepper
-from .spectral import derivative
-from .utils import (
+from . import _metrics as metrics
+from . import _poisson as poisson
+from . import ic, nonlin_fun, normalized, stepper
+from ._forced_stepper import ForcedStepper
+from ._repeated_stepper import RepeatedStepper
+from ._spectral import derivative
+from ._utils import (
     build_ic_set,
     get_animation,
     get_grid,

exponax/base_stepper.py → exponax/_base_stepper.py

@@ -2,14 +2,14 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex, Float
 
-from .exponential_integrators import ETDRK0, ETDRK1, ETDRK2, ETDRK3, ETDRK4, BaseETDRK
-from .nonlin_fun import BaseNonlinearFun
-from .spectral import (
+from ._exponential_integrators import ETDRK0, ETDRK1, ETDRK2, ETDRK3, ETDRK4, BaseETDRK
+from ._spectral import (
     build_derivative_operator,
     space_indices,
     spatial_shape,
     wavenumber_shape,
 )
+from .nonlin_fun import BaseNonlinearFun
 
 
 class BaseStepper(eqx.Module):

exponax/forced_stepper.py → exponax/_forced_stepper.py

@@ -1,7 +1,7 @@
 import equinox as eqx
 from jaxtyping import Array, Complex, Float
 
-from .base_stepper import BaseStepper
+from ._base_stepper import BaseStepper
 
 
 class ForcedStepper(eqx.Module):

exponax/poisson.py → exponax/_poisson.py

@@ -2,7 +2,7 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex, Float
 
-from .spectral import build_derivative_operator, build_laplace_operator, spatial_shape
+from ._spectral import build_derivative_operator, build_laplace_operator, spatial_shape
 
 
 class Poisson(eqx.Module):

exponax/repeated_stepper.py → exponax/_repeated_stepper.py

@@ -1,8 +1,8 @@
 import equinox as eqx
 from jaxtyping import Array, Complex, Float
 
-from .base_stepper import BaseStepper
-from .utils import repeat
+from ._base_stepper import BaseStepper
+from ._utils import repeat
 
 
 class RepeatedStepper(eqx.Module):

exponax/ic/__init__.py

@@ -1,8 +1,8 @@
-from .base_ic import BaseIC, BaseRandomICGenerator
-from .diffused_noise import DiffusedNoise
-from .gaussian_random_field import GaussianRandomField
-from .multi_channel import MultiChannelIC, RandomMultiChannelICGenerator
-from .truncated_fourier_series import RandomTruncatedFourierSeries
+from ._base_ic import BaseIC, BaseRandomICGenerator
+from ._diffused_noise import DiffusedNoise
+from ._gaussian_random_field import GaussianRandomField
+from ._multi_channel import MultiChannelIC, RandomMultiChannelICGenerator
+from ._truncated_fourier_series import RandomTruncatedFourierSeries
 
 __all__ = [
     "BaseIC",

exponax/ic/base_ic.py → exponax/ic/_base_ic.py

@@ -3,7 +3,7 @@
 import equinox as eqx
 from jaxtyping import Array, Float, PRNGKeyArray
 
-from ..utils import get_grid
+from .._utils import get_grid
 
 
 class BaseIC(eqx.Module, ABC):

exponax/ic/diffused_noise.py → exponax/ic/_diffused_noise.py

@@ -2,9 +2,9 @@
 import jax.random as jr
 from jaxtyping import Array, Float, PRNGKeyArray
 
-from ..spectral import spatial_shape
+from .._spectral import spatial_shape
 from ..stepper import Diffusion
-from .base_ic import BaseRandomICGenerator
+from ._base_ic import BaseRandomICGenerator
 
 
 class DiffusedNoise(BaseRandomICGenerator):

exponax/ic/gaussian_random_field.py → exponax/ic/_gaussian_random_field.py

@@ -2,13 +2,13 @@
 import jax.random as jr
 from jaxtyping import Array, Float, PRNGKeyArray
 
-from ..spectral import (
+from .._spectral import (
     build_scaled_wavenumbers,
     space_indices,
     spatial_shape,
     wavenumber_shape,
 )
-from .base_ic import BaseRandomICGenerator
+from ._base_ic import BaseRandomICGenerator
 
 
 class GaussianRandomField(BaseRandomICGenerator):

exponax/ic/multi_channel.py → exponax/ic/_multi_channel.py

@@ -4,7 +4,7 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Float, PRNGKeyArray
 
-from .base_ic import BaseIC, BaseRandomICGenerator
+from ._base_ic import BaseIC, BaseRandomICGenerator
 
 
 class MultiChannelIC(eqx.Module):

exponax/ic/truncated_fourier_series.py → exponax/ic/_truncated_fourier_series.py

@@ -2,14 +2,14 @@
 import jax.random as jr
 from jaxtyping import Array, Float, PRNGKeyArray
 
-from ..spectral import (
+from .._spectral import (
     build_scaling_array,
     low_pass_filter_mask,
     space_indices,
     spatial_shape,
     wavenumber_shape,
 )
-from .base_ic import BaseRandomICGenerator
+from ._base_ic import BaseRandomICGenerator
 
 
 class RandomTruncatedFourierSeries(BaseRandomICGenerator):

exponax/nonlin_fun/__init__.py

@@ -1,14 +1,17 @@
-from .base import BaseNonlinearFun
-from .convection import ConvectionNonlinearFun
-from .gradient_norm import GradientNormNonlinearFun
-from .polynomial import PolynomialNonlinearFun
-from .reaction import (
+from ._base import BaseNonlinearFun
+from ._convection import ConvectionNonlinearFun
+from ._gradient_norm import GradientNormNonlinearFun
+from ._polynomial import PolynomialNonlinearFun
+from ._reaction import (
     BelousovZhabotinskyNonlinearFun,
     CahnHilliardNonlinearFun,
     GrayScottNonlinearFun,
 )
-from .vorticity_convection import VorticityConvection2d, VorticityConvection2dKolmogorov
-from .zero import ZeroNonlinearFun
+from ._vorticity_convection import (
+    VorticityConvection2d,
+    VorticityConvection2dKolmogorov,
+)
+from ._zero import ZeroNonlinearFun
 
 __all__ = [
     "BaseNonlinearFun",

exponax/nonlin_fun/base.py → exponax/nonlin_fun/_base.py

@@ -3,7 +3,7 @@
 import equinox as eqx
 from jaxtyping import Array, Bool, Complex
 
-from ..spectral import low_pass_filter_mask, wavenumber_shape
+from .._spectral import low_pass_filter_mask, wavenumber_shape
 
 
 class BaseNonlinearFun(eqx.Module, ABC):

exponax/nonlin_fun/convection.py → exponax/nonlin_fun/_convection.py

@@ -2,8 +2,8 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex, Float
 
-from ..spectral import space_indices, spatial_shape
-from .base import BaseNonlinearFun
+from .._spectral import space_indices, spatial_shape
+from ._base import BaseNonlinearFun
 
 
 class ConvectionNonlinearFun(BaseNonlinearFun):

exponax/nonlin_fun/gradient_norm.py → exponax/nonlin_fun/_gradient_norm.py

@@ -2,8 +2,8 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex, Float
 
-from ..spectral import space_indices, spatial_shape
-from .base import BaseNonlinearFun
+from .._spectral import space_indices, spatial_shape
+from ._base import BaseNonlinearFun
 
 
 class GradientNormNonlinearFun(BaseNonlinearFun):

exponax/nonlin_fun/polynomial.py → exponax/nonlin_fun/_polynomial.py

@@ -1,8 +1,8 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex
 
-from ..spectral import space_indices, spatial_shape
-from .base import BaseNonlinearFun
+from .._spectral import space_indices, spatial_shape
+from ._base import BaseNonlinearFun
 
 
 class PolynomialNonlinearFun(BaseNonlinearFun):

exponax/nonlin_fun/reaction.py → exponax/nonlin_fun/_reaction.py

@@ -5,8 +5,8 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex
 
-from ..spectral import build_laplace_operator, space_indices, spatial_shape
-from .base import BaseNonlinearFun
+from .._spectral import build_laplace_operator, space_indices, spatial_shape
+from ._base import BaseNonlinearFun
 
 
 class GrayScottNonlinearFun(BaseNonlinearFun):

exponax/nonlin_fun/vorticity_convection.py → exponax/nonlin_fun/_vorticity_convection.py

@@ -1,8 +1,8 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex
 
-from ..spectral import build_laplace_operator, build_scaling_array, build_wavenumbers
-from .base import BaseNonlinearFun
+from .._spectral import build_laplace_operator, build_scaling_array, build_wavenumbers
+from ._base import BaseNonlinearFun
 
 
 class VorticityConvection2d(BaseNonlinearFun):

exponax/nonlin_fun/zero.py → exponax/nonlin_fun/_zero.py

@@ -1,7 +1,7 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex
 
-from .base import BaseNonlinearFun
+from ._base import BaseNonlinearFun
 
 
 class ZeroNonlinearFun(BaseNonlinearFun):

exponax/normalized/__init__.py

@@ -1,7 +1,7 @@
-from .convection import NormalizedConvectionStepper
-from .gradient_norm import NormalizedGradientNormStepper
-from .linear import NormalizedLinearStepper
-from .utils import (
+from ._convection import NormalizedConvectionStepper
+from ._gradient_norm import NormalizedGradientNormStepper
+from ._linear import NormalizedLinearStepper
+from ._utils import (
     denormalize_coefficients,
     denormalize_convection_scale,
     denormalize_gradient_norm_scale,

exponax/normalized/convection.py → exponax/normalized/_convection.py

@@ -1,7 +1,7 @@
 import jax.numpy as jnp
 from jaxtyping import Array
 
-from ..base_stepper import BaseStepper
+from .._base_stepper import BaseStepper
 from ..nonlin_fun import ConvectionNonlinearFun
 
 

exponax/normalized/gradient_norm.py → exponax/normalized/_gradient_norm.py

@@ -1,7 +1,7 @@
 import jax.numpy as jnp
 from jaxtyping import Array
 
-from ..base_stepper import BaseStepper
+from .._base_stepper import BaseStepper
 from ..nonlin_fun import GradientNormNonlinearFun
 
 

exponax/normalized/linear.py → exponax/normalized/_linear.py

@@ -1,7 +1,7 @@
 import jax.numpy as jnp
 from jaxtyping import Array
 
-from ..base_stepper import BaseStepper
+from .._base_stepper import BaseStepper
 from ..nonlin_fun import ZeroNonlinearFun
 
 

exponax/stepper/__init__.py

@@ -1,19 +1,19 @@
-from .burgers import Burgers
-from .convection import GeneralConvectionStepper
-from .gradient_norm import GeneralGradientNormStepper
-from .korteveg_de_vries import KortewegDeVries
-from .kuramoto_sivashinsky import KuramotoSivashinsky, KuramotoSivashinskyConservative
-from .linear import (
+from ._burgers import Burgers
+from ._convection import GeneralConvectionStepper
+from ._gradient_norm import GeneralGradientNormStepper
+from ._korteveg_de_vries import KortewegDeVries
+from ._kuramoto_sivashinsky import KuramotoSivashinsky, KuramotoSivashinskyConservative
+from ._linear import (
     Advection,
     AdvectionDiffusion,
     Diffusion,
     Dispersion,
     GeneralLinearStepper,
     HyperDiffusion,
 )
-from .navier_stokes import KolmogorovFlowVorticity2d, NavierStokesVorticity2d
-from .nikolaevskiy import Nikolaevskiy, NikolaevskiyConservative
-from .reaction import (
+from ._navier_stokes import KolmogorovFlowVorticity2d, NavierStokesVorticity2d
+from ._nikolaevskiy import Nikolaevskiy, NikolaevskiyConservative
+from ._reaction import (
     AllenCahn,
     BelousovZhabotinsky,
     CahnHilliard,

exponax/stepper/burgers.py → exponax/stepper/_burgers.py

@@ -1,8 +1,8 @@
 from jaxtyping import Array, Complex
 
-from ..base_stepper import BaseStepper
+from .._base_stepper import BaseStepper
+from .._spectral import build_laplace_operator
 from ..nonlin_fun import ConvectionNonlinearFun
-from ..spectral import build_laplace_operator
 
 
 class Burgers(BaseStepper):

exponax/stepper/convection.py → exponax/stepper/_convection.py

@@ -1,7 +1,7 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex
 
-from ..base_stepper import BaseStepper
+from .._base_stepper import BaseStepper
 from ..nonlin_fun import ConvectionNonlinearFun
 
 

exponax/stepper/gradient_norm.py → exponax/stepper/_gradient_norm.py

@@ -1,7 +1,7 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex
 
-from ..base_stepper import BaseStepper
+from .._base_stepper import BaseStepper
 from ..nonlin_fun import GradientNormNonlinearFun
 
 

exponax/stepper/korteveg_de_vries.py → exponax/stepper/_korteveg_de_vries.py

@@ -3,9 +3,9 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex, Float
 
-from ..base_stepper import BaseStepper
+from .._base_stepper import BaseStepper
+from .._spectral import build_gradient_inner_product_operator, build_laplace_operator
 from ..nonlin_fun import ConvectionNonlinearFun
-from ..spectral import build_gradient_inner_product_operator, build_laplace_operator
 
 D = TypeVar("D")
 

exponax/stepper/kuramoto_sivashinsky.py → exponax/stepper/_kuramoto_sivashinsky.py

@@ -1,8 +1,8 @@
 from jaxtyping import Array, Complex
 
-from ..base_stepper import BaseStepper
+from .._base_stepper import BaseStepper
+from .._spectral import build_laplace_operator
 from ..nonlin_fun import ConvectionNonlinearFun, GradientNormNonlinearFun
-from ..spectral import build_laplace_operator
 
 
 class KuramotoSivashinsky(BaseStepper):

exponax/stepper/linear.py → exponax/stepper/_linear.py

@@ -3,9 +3,9 @@
 import jax.numpy as jnp
 from jaxtyping import Array, Complex, Float
 
-from ..base_stepper import BaseStepper
+from .._base_stepper import BaseStepper
+from .._spectral import build_gradient_inner_product_operator, build_laplace_operator
 from ..nonlin_fun import ZeroNonlinearFun
-from ..spectral import build_gradient_inner_product_operator, build_laplace_operator
 
 D = TypeVar("D")
 

exponax/stepper/navier_stokes.py → exponax/stepper/_navier_stokes.py

@@ -1,8 +1,8 @@
 from jaxtyping import Array, Complex
 
-from ..base_stepper import BaseStepper
+from .._base_stepper import BaseStepper
+from .._spectral import build_laplace_operator
 from ..nonlin_fun import VorticityConvection2d, VorticityConvection2dKolmogorov
-from ..spectral import build_laplace_operator
 
 
 class NavierStokesVorticity2d(BaseStepper):

exponax/stepper/nikolaevskiy.py → exponax/stepper/_nikolaevskiy.py

@@ -1,8 +1,8 @@
 from jaxtyping import Array, Complex
 
-from ..base_stepper import BaseStepper
+from .._base_stepper import BaseStepper
+from .._spectral import build_laplace_operator
 from ..nonlin_fun import ConvectionNonlinearFun, GradientNormNonlinearFun
-from ..spectral import build_laplace_operator
 
 
 class Nikolaevskiy(BaseStepper):