better module init

wcxve · Mar 2, 2024 · 46ce257 · 46ce257
1 parent 5940592
commit 46ce257
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 27 deletions.
diff --git a/src/elisa/model/__init__.py b/src/elisa/model/__init__.py
@@ -1,13 +1,44 @@
-from . import add, conv, model, mul, parameter
+from . import add, conv, mul
 from .add import *  # noqa: F403
 from .conv import *  # noqa: F403
-from .model import *  # noqa: F403
+from .model import (
+    AdditiveComponent,
+    AnaIntAdditive,
+    AnaIntMultiplicative,
+    ConvolutionComponent,
+    MultiplicativeComponent,
+    NumIntAdditive,
+    NumIntMultiplicative,
+)
 from .mul import *  # noqa: F403
-from .parameter import *  # noqa: F403
+from .parameter import (
+    CompositeParameter,
+    ConstantInterval,
+    ConstantValue,
+    Parameter,
+    ParameterBase,
+    UniformParameter,
+)
 
 __all__ = (
-    model.__all__
-    + parameter.__all__
+    [
+        'AdditiveComponent',
+        'MultiplicativeComponent',
+        'ConvolutionComponent',
+        'AnaIntAdditive',
+        'NumIntAdditive',
+        'AnaIntMultiplicative',
+        'NumIntMultiplicative',
+    ]
+    + [
+        'ParameterBase',
+        'Parameter',
+        'UniformParameter',
+        'ConstantValue',
+        'ConstantInterval',
+        'CompositeParameter',
+        # 'GPParameter',
+    ]
     + add.__all__
     + mul.__all__
     + conv.__all__

diff --git a/src/elisa/model/model.py b/src/elisa/model/model.py
@@ -37,16 +37,6 @@
     T,
 )
 
-__all__ = [
-    'AdditiveComponent',
-    'MultiplicativeComponent',
-    'ConvolutionComponent',
-    'AnaIntAdditive',
-    'NumIntAdditive',
-    'AnaIntMultiplicative',
-    'NumIntMultiplicative',
-]
-
 
 class ModelBase(metaclass=ABCMeta):
     """Base model class."""
@@ -583,6 +573,12 @@ def flux(
     ) -> jax.Array | dict[str, jax.Array]:
         r"""Calculate the flux of model between `emin` and `emax`.
 
+        Warnings
+        --------
+        The flux is calculated by trapzoidal rule, which may not be accurate
+        if not enough energy bins are used when the difference between
+        `emin` and emax` is large.
+
         Parameters
         ----------
         emin : float or int

diff --git a/src/elisa/model/parameter.py b/src/elisa/model/parameter.py
@@ -8,14 +8,13 @@
 import jax.numpy as jnp
 from numpyro.distributions import Distribution, LogUniform, Uniform
 
+# from tinygp import kernels, means, noise
 from elisa.util.integrate import (
     AdaptQuadMethod,
     IntegralFactory,
     make_integral_factory,
 )
 from elisa.util.misc import build_namespace
-
-# from tinygp import kernels, means, noise
 from elisa.util.typing import (
     CompID,
     CompParamName,
@@ -26,16 +25,6 @@
     ParamIDValMapping,
 )
 
-__all__ = [
-    'ParameterBase',
-    'Parameter',
-    'UniformParameter',
-    'ConstantValue',
-    'ConstantInterval',
-    'CompositeParameter',
-    # 'GPParameter',
-]
-
 
 class AssignmentTracker:
     """Track component assignment of a parameter."""