From 46ce257706f98fcb43099ea25be9db88f8c246c1 Mon Sep 17 00:00:00 2001
From: xuewc <wcxuemail@gmail.com>
Date: Sat, 2 Mar 2024 13:58:30 +0800
Subject: [PATCH] better module init

---
 src/elisa/model/__init__.py  | 41 +++++++++++++++++++++++++++++++-----
 src/elisa/model/model.py     | 16 ++++++--------
 src/elisa/model/parameter.py | 13 +-----------
 3 files changed, 43 insertions(+), 27 deletions(-)

diff --git a/src/elisa/model/__init__.py b/src/elisa/model/__init__.py
index f81b0e7a..6020e3ca 100644
--- 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
index 171a6a3d..921d43ae 100644
--- 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
index ed843aed..e6e870cc 100644
--- 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."""