From 60e87db1def95a6b8c40827323ae4f7029132dba Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 21 Dec 2023 16:24:51 +0100 Subject: [PATCH 001/346] format --- src/hapsira/core/_jit.py | 157 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 157 insertions(+) create mode 100644 src/hapsira/core/_jit.py diff --git a/src/hapsira/core/_jit.py b/src/hapsira/core/_jit.py new file mode 100644 index 000000000..f1bdb9c2c --- /dev/null +++ b/src/hapsira/core/_jit.py @@ -0,0 +1,157 @@ +import os + +import numba as nb + +TARGET = os.environ.get("HAPSIRA_TARGET", "cpu") + +if TARGET not in ( + "cpu", + "parallel", + "cuda", +): # numba 0.54.0, 19 August 2021, removed AMD ROCm target + raise ValueError(f'unknown target "{TARGET:s}"') +if TARGET == "cuda": + from numba import ( + cuda, + ) + +INLINE = os.environ.get( + "HAPSIRA_INLINE", "never" +) # currently only relevant for helpers on cpu and parallel targets +if INLINE not in ("always", "never"): + raise ValueError(f'unknown value for inline "{INLINE:s}"') + +PRECISIONS = ("f4", "f8") # TODO allow f2, i.e. half, for CUDA at least? + +# cuda.jit does not allow multiple signatures, i.e. stuff can only be compiled +# for one precision level, see https://github.com/numba/numba/issues/3226 +CUDA_PRECISION = os.environ.get("HAPSIRA_CUDA_PRECISION", "f8") +if CUDA_PRECISION not in PRECISIONS: + raise ValueError(f'unknown floating point precision "{CUDA_PRECISION:s}"') + +NOPYTHON = True # only for debugging, True by default + + +def _parse_signatures(signature): + """ + Automatically generate signatures for single and double + """ + if not any( + notation in signature for notation in ("f", "V") + ): # leave this signature as it is + return signature + if any(level in signature for level in PRECISIONS): # leave this signature as it is + return signature + signature = signature.replace( + "V", "Tuple([f,f,f])" + ) # TODO hope for support of "f[:]" return values in cuda target + if TARGET == "cuda": + return signature.replace("f", CUDA_PRECISION) + return [signature.replace("f", dtype) for dtype in PRECISIONS] + + +def hjit(*args, **kwargs): + """ + Scalar helper, pre-configured, internal, switches compiler targets. + Functions decorated by it can only be called directly if TARGET is cpu or parallel. + """ + + if len(args) == 1 and callable(args[0]): + func = args[0] + args = tuple() + else: + func = None + + if len(args) > 0 and isinstance(args[0], str): + args = _parse_signatures(args[0]), *args[1:] + + def wrapper(func): + cfg = {} + if TARGET in ("cpu", "parallel"): + cfg.update({"nopython": NOPYTHON, "inline": INLINE}) + if TARGET == "cuda": + cfg.update({"device": True, "inline": True}) + cfg.update(kwargs) + + wjit = cuda.jit if TARGET == "cuda" else nb.jit + + return wjit( + *args, + **cfg, + )(func) + + if func is not None: + return wrapper(func) + + return wrapper + + +def vjit(*args, **kwargs): + """ + Vectorize on array, pre-configured, user-facing, switches compiler targets. + Functions decorated by it can always be called directly if needed. + """ + + if len(args) == 1 and callable(args[0]): + func = args[0] + args = tuple() + else: + func = None + + if len(args) > 0 and isinstance(args[0], str): + args = _parse_signatures(args[0]), *args[1:] + + def wrapper(func): + cfg = {"target": TARGET} + if TARGET in ("cpu", "parallel"): + cfg.update({"nopython": NOPYTHON}) + cfg.update(kwargs) + + return nb.vectorize( + *args, + **cfg, + )(func) + + if func is not None: + return wrapper(func) + + return wrapper + + +def jit(*args, **kwargs): + """ + Regular (n)jit, pre-configured, potentially user-facing, always CPU compiler target. + Functions decorated by it can only be called directly. + """ + + if len(args) == 1 and callable(args[0]): + func = args[0] + args = tuple() + else: + func = None + + def wrapper(func): + cfg = {"nopython": NOPYTHON, "inline": "never"} # inline in ('always', 'never') + cfg.update(kwargs) + + return nb.jit( + *args, + **cfg, + )(func) + + if func is not None: + return wrapper(func) + + return wrapper + + +__all__ = [ + "CUDA_PRECISION", + "INLINE", + "NOPYTHON", + "PRECISIONS", + "TARGET", + "hjit", + "jit", + "vjit", +] From 5d0ce6eee988bdadfd0e4e79900726ef01eab904 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 22 Dec 2023 16:35:27 +0100 Subject: [PATCH 002/346] const sub black --- src/hapsira/const.py | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 src/hapsira/const.py diff --git a/src/hapsira/const.py b/src/hapsira/const.py new file mode 100644 index 000000000..1bb32710e --- /dev/null +++ b/src/hapsira/const.py @@ -0,0 +1,2 @@ +TRUE = ("true", "1", "yes") +FALSE = ("false", "0", "no") From 771f310b1b2ba7cefd0ee562d4adb9167c534fd5 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 22 Dec 2023 16:35:51 +0100 Subject: [PATCH 003/346] error sub --- src/hapsira/errors.py | 2 ++ 1 file changed, 2 insertions(+) create mode 100644 src/hapsira/errors.py diff --git a/src/hapsira/errors.py b/src/hapsira/errors.py new file mode 100644 index 000000000..838069523 --- /dev/null +++ b/src/hapsira/errors.py @@ -0,0 +1,2 @@ +class JitError(Exception): + pass From c5c5b302f2d48563b372faf5e49c279caca4cd0a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 22 Dec 2023 16:37:34 +0100 Subject: [PATCH 004/346] jit draft, not working --- src/hapsira/core/_jit.py | 64 +++++++++++++++++++++++++++++----------- 1 file changed, 47 insertions(+), 17 deletions(-) diff --git a/src/hapsira/core/_jit.py b/src/hapsira/core/_jit.py index f1bdb9c2c..6d40b8f71 100644 --- a/src/hapsira/core/_jit.py +++ b/src/hapsira/core/_jit.py @@ -1,24 +1,52 @@ +from enum import Enum, auto import os import numba as nb +from poliastro.errors import JitError -TARGET = os.environ.get("HAPSIRA_TARGET", "cpu") -if TARGET not in ( - "cpu", - "parallel", - "cuda", -): # numba 0.54.0, 19 August 2021, removed AMD ROCm target - raise ValueError(f'unknown target "{TARGET:s}"') +class TARGETS(Enum): + """ + JIT targets + numba 0.54.0, 19 August 2021, removed AMD ROCm target + """ + + cpu = auto() + parallel = auto() + cuda = auto() + + @classmethod + def get_default(cls): + """ + default JIT target + """ + + return cls.cpu + + +TARGET = os.environ.get("HAPSIRA_TARGET", TARGETS.default().name) + +try: + TARGET = TARGETS[TARGET] +except KeyError as e: + raise JitError( + f'unknown target "{TARGET:s}"; known targets are {repr(TARGETS):s}' + ) from e + if TARGET == "cuda": - from numba import ( - cuda, - ) + from numba import cuda + + if not cuda.is_available(): + raise JitError('no GPU for selected target "cuda" found') INLINE = os.environ.get( - "HAPSIRA_INLINE", "never" + "HAPSIRA_INLINE", "0" ) # currently only relevant for helpers on cpu and parallel targets -if INLINE not in ("always", "never"): +if INLINE.lower() in ("true", "1", "yes"): + INLINE = True +elif INLINE.lower() in ("false", "0", "no"): + INLINE = False +else: raise ValueError(f'unknown value for inline "{INLINE:s}"') PRECISIONS = ("f4", "f8") # TODO allow f2, i.e. half, for CUDA at least? @@ -68,7 +96,9 @@ def hjit(*args, **kwargs): def wrapper(func): cfg = {} if TARGET in ("cpu", "parallel"): - cfg.update({"nopython": NOPYTHON, "inline": INLINE}) + cfg.update( + {"nopython": NOPYTHON, "inline": "always" if INLINE else "never"} + ) if TARGET == "cuda": cfg.update({"device": True, "inline": True}) cfg.update(kwargs) @@ -118,10 +148,10 @@ def wrapper(func): return wrapper -def jit(*args, **kwargs): +def sjit(*args, **kwargs): """ - Regular (n)jit, pre-configured, potentially user-facing, always CPU compiler target. - Functions decorated by it can only be called directly. + Regular "scalar" (n)jit, pre-configured, potentially user-facing, always CPU compiler target. + Functions decorated by it can always be called directly if needed. """ if len(args) == 1 and callable(args[0]): @@ -152,6 +182,6 @@ def wrapper(func): "PRECISIONS", "TARGET", "hjit", - "jit", "vjit", + "sjit", ] From 0a508a13cd357101f0bc92f17f431bd397c0e848 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 24 Dec 2023 17:37:51 +0100 Subject: [PATCH 005/346] jit cleanup; structure --- src/hapsira/core/_jit.py | 175 ++++++++++++++++++++++++--------------- 1 file changed, 108 insertions(+), 67 deletions(-) diff --git a/src/hapsira/core/_jit.py b/src/hapsira/core/_jit.py index 6d40b8f71..1b21e9d05 100644 --- a/src/hapsira/core/_jit.py +++ b/src/hapsira/core/_jit.py @@ -1,186 +1,227 @@ from enum import Enum, auto import os +from typing import Callable import numba as nb +from numba import cuda from poliastro.errors import JitError +def _str2bool(value: str) -> bool: + """ + Helper for parsing environment variables + """ + + if value.strip().lower() in ("true", "1", "yes"): + return True + if value.strip().lower() in ("false", "0", "no"): + return False + raise ValueError(f'can not convert value "{value:s}" to bool') + + class TARGETS(Enum): """ JIT targets - numba 0.54.0, 19 August 2021, removed AMD ROCm target """ cpu = auto() parallel = auto() cuda = auto() + # numba 0.54.0, 19 August 2021, removed AMD ROCm target @classmethod def get_default(cls): """ - default JIT target + Default JIT target """ return cls.cpu + @classmethod + def get_current(cls): + """ + Current JIT target + """ + + name = os.environ.get("HAPSIRA_TARGET", None) -TARGET = os.environ.get("HAPSIRA_TARGET", TARGETS.default().name) + if name is None: + target = cls.get_default() + else: + try: + target = cls[name] + except KeyError as e: + raise JitError( + f'unknown target "{name:s}"; known targets are {repr(cls):s}' + ) from e -try: - TARGET = TARGETS[TARGET] -except KeyError as e: - raise JitError( - f'unknown target "{TARGET:s}"; known targets are {repr(TARGETS):s}' - ) from e + if target is cls.cuda and not cuda.is_available(): + raise JitError('selected target "cuda" is not available') -if TARGET == "cuda": - from numba import cuda + return target - if not cuda.is_available(): - raise JitError('no GPU for selected target "cuda" found') -INLINE = os.environ.get( - "HAPSIRA_INLINE", "0" -) # currently only relevant for helpers on cpu and parallel targets -if INLINE.lower() in ("true", "1", "yes"): - INLINE = True -elif INLINE.lower() in ("false", "0", "no"): - INLINE = False -else: - raise ValueError(f'unknown value for inline "{INLINE:s}"') +TARGET = TARGETS.get_current() -PRECISIONS = ("f4", "f8") # TODO allow f2, i.e. half, for CUDA at least? +INLINE = _str2bool( + os.environ.get("HAPSIRA_INLINE", "1" if TARGET is TARGET.cuda else "0") +) # currently only relevant for helpers on cpu and parallel targets -# cuda.jit does not allow multiple signatures, i.e. stuff can only be compiled -# for one precision level, see https://github.com/numba/numba/issues/3226 -CUDA_PRECISION = os.environ.get("HAPSIRA_CUDA_PRECISION", "f8") -if CUDA_PRECISION not in PRECISIONS: - raise ValueError(f'unknown floating point precision "{CUDA_PRECISION:s}"') +NOPYTHON = _str2bool( + os.environ.get("HAPSIRA_NOPYTHON", "1") +) # only for debugging, True by default -NOPYTHON = True # only for debugging, True by default +PRECISIONS = ("f4", "f8") # TODO allow f2, i.e. half, for CUDA at least? -def _parse_signatures(signature): +def _parse_signatures(signature: str) -> str | list[str]: """ Automatically generate signatures for single and double """ + if not any( notation in signature for notation in ("f", "V") ): # leave this signature as it is + # TODO warn! return signature + if any(level in signature for level in PRECISIONS): # leave this signature as it is + # TODO warn! return signature + signature = signature.replace( "V", "Tuple([f,f,f])" - ) # TODO hope for support of "f[:]" return values in cuda target - if TARGET == "cuda": - return signature.replace("f", CUDA_PRECISION) + ) # TODO hope for support of "f[:]" return values in cuda target; 2D/4D vectors? return [signature.replace("f", dtype) for dtype in PRECISIONS] -def hjit(*args, **kwargs): +def hjit(*args, **kwargs) -> Callable: """ Scalar helper, pre-configured, internal, switches compiler targets. Functions decorated by it can only be called directly if TARGET is cpu or parallel. """ if len(args) == 1 and callable(args[0]): - func = args[0] + outer_func = args[0] args = tuple() else: - func = None + outer_func = None if len(args) > 0 and isinstance(args[0], str): args = _parse_signatures(args[0]), *args[1:] - def wrapper(func): - cfg = {} - if TARGET in ("cpu", "parallel"): - cfg.update( - {"nopython": NOPYTHON, "inline": "always" if INLINE else "never"} + def wrapper(inner_func: Callable) -> Callable: + """ + Applies JIT + """ + + if TARGET in (TARGETS.cpu, TARGETS.parallel): + wjit = nb.jit + cfg = dict( + nopython=NOPYTHON, + inline="always" if INLINE else "never", + ) + elif TARGET is TARGETS.cuda: + wjit = cuda.jit + cfg = dict( + device=True, + inline=INLINE, + ) + else: + JitError( + f'unknown target "{repr(TARGET):s}"; known targets are {repr(TARGETS):s}' ) - if TARGET == "cuda": - cfg.update({"device": True, "inline": True}) cfg.update(kwargs) - wjit = cuda.jit if TARGET == "cuda" else nb.jit - return wjit( *args, **cfg, - )(func) + )(inner_func) - if func is not None: - return wrapper(func) + if outer_func is not None: + return wrapper(outer_func) return wrapper -def vjit(*args, **kwargs): +def vjit(*args, **kwargs) -> Callable: """ Vectorize on array, pre-configured, user-facing, switches compiler targets. Functions decorated by it can always be called directly if needed. """ if len(args) == 1 and callable(args[0]): - func = args[0] + outer_func = args[0] args = tuple() else: - func = None + outer_func = None if len(args) > 0 and isinstance(args[0], str): args = _parse_signatures(args[0]), *args[1:] - def wrapper(func): - cfg = {"target": TARGET} - if TARGET in ("cpu", "parallel"): - cfg.update({"nopython": NOPYTHON}) + def wrapper(inner_func: Callable) -> Callable: + """ + Applies JIT + """ + + cfg = dict( + target=TARGET.name, + ) + if TARGET is not TARGETS.cuda: + cfg["nopython"] = NOPYTHON cfg.update(kwargs) return nb.vectorize( *args, **cfg, - )(func) + )(inner_func) - if func is not None: - return wrapper(func) + if outer_func is not None: + return wrapper(outer_func) return wrapper -def sjit(*args, **kwargs): +def sjit(*args, **kwargs) -> Callable: """ Regular "scalar" (n)jit, pre-configured, potentially user-facing, always CPU compiler target. Functions decorated by it can always be called directly if needed. """ if len(args) == 1 and callable(args[0]): - func = args[0] + outer_func = args[0] args = tuple() else: - func = None + pass + + def wrapper(inner_func: Callable) -> Callable: + """ + Applies JIT + """ - def wrapper(func): - cfg = {"nopython": NOPYTHON, "inline": "never"} # inline in ('always', 'never') + cfg = dict( + nopython=NOPYTHON, + inline="always" if INLINE else "never", + ) cfg.update(kwargs) return nb.jit( *args, **cfg, - )(func) + )(inner_func) - if func is not None: - return wrapper(func) + if outer_func is not None: + return wrapper(outer_func) return wrapper __all__ = [ - "CUDA_PRECISION", "INLINE", "NOPYTHON", "PRECISIONS", "TARGET", + "TARGETS", "hjit", "vjit", "sjit", From 8b0205b08d2396f82d35ef89e32637b882c4fcc4 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 24 Dec 2023 17:38:08 +0100 Subject: [PATCH 006/346] rm consts --- src/hapsira/const.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/hapsira/const.py b/src/hapsira/const.py index 1bb32710e..e69de29bb 100644 --- a/src/hapsira/const.py +++ b/src/hapsira/const.py @@ -1,2 +0,0 @@ -TRUE = ("true", "1", "yes") -FALSE = ("false", "0", "no") From 5df0c8c7b807eaeab5d538a887ff40c6e9acaa41 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 26 Dec 2023 22:59:28 +0100 Subject: [PATCH 007/346] rm const sub --- src/hapsira/const.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 src/hapsira/const.py diff --git a/src/hapsira/const.py b/src/hapsira/const.py deleted file mode 100644 index e69de29bb..000000000 From 39054307e7b7206ba6706fbe77170f6296d54bde Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 27 Dec 2023 11:18:14 +0100 Subject: [PATCH 008/346] moved math into core --- src/hapsira/{ => core}/_math/__init__.py | 0 src/hapsira/{ => core}/_math/integrate.py | 0 src/hapsira/{ => core}/_math/interpolate.py | 0 src/hapsira/{ => core}/_math/ivp.py | 0 src/hapsira/{ => core}/_math/linalg.py | 0 src/hapsira/{ => core}/_math/optimize.py | 0 src/hapsira/{ => core}/_math/special.py | 0 src/hapsira/core/czml_utils.py | 2 +- src/hapsira/core/elements.py | 3 ++- src/hapsira/core/events.py | 3 ++- src/hapsira/core/flybys.py | 2 +- src/hapsira/core/iod.py | 4 ++-- src/hapsira/core/maneuver.py | 3 ++- src/hapsira/core/perturbations.py | 3 ++- src/hapsira/core/propagation/cowell.py | 3 ++- src/hapsira/core/propagation/vallado.py | 4 ++-- src/hapsira/core/spheroid_location.py | 2 +- src/hapsira/core/thrust/change_a_inc.py | 3 ++- src/hapsira/core/thrust/change_argp.py | 3 ++- src/hapsira/core/thrust/change_ecc_inc.py | 3 ++- src/hapsira/earth/atmosphere/coesa62.py | 2 +- src/hapsira/ephem.py | 2 +- src/hapsira/threebody/restricted.py | 2 +- src/hapsira/twobody/events.py | 2 +- src/hapsira/util.py | 2 +- tests/test_hyper.py | 2 +- tests/test_stumpff.py | 2 +- 27 files changed, 30 insertions(+), 22 deletions(-) rename src/hapsira/{ => core}/_math/__init__.py (100%) rename src/hapsira/{ => core}/_math/integrate.py (100%) rename src/hapsira/{ => core}/_math/interpolate.py (100%) rename src/hapsira/{ => core}/_math/ivp.py (100%) rename src/hapsira/{ => core}/_math/linalg.py (100%) rename src/hapsira/{ => core}/_math/optimize.py (100%) rename src/hapsira/{ => core}/_math/special.py (100%) diff --git a/src/hapsira/_math/__init__.py b/src/hapsira/core/_math/__init__.py similarity index 100% rename from src/hapsira/_math/__init__.py rename to src/hapsira/core/_math/__init__.py diff --git a/src/hapsira/_math/integrate.py b/src/hapsira/core/_math/integrate.py similarity index 100% rename from src/hapsira/_math/integrate.py rename to src/hapsira/core/_math/integrate.py diff --git a/src/hapsira/_math/interpolate.py b/src/hapsira/core/_math/interpolate.py similarity index 100% rename from src/hapsira/_math/interpolate.py rename to src/hapsira/core/_math/interpolate.py diff --git a/src/hapsira/_math/ivp.py b/src/hapsira/core/_math/ivp.py similarity index 100% rename from src/hapsira/_math/ivp.py rename to src/hapsira/core/_math/ivp.py diff --git a/src/hapsira/_math/linalg.py b/src/hapsira/core/_math/linalg.py similarity index 100% rename from src/hapsira/_math/linalg.py rename to src/hapsira/core/_math/linalg.py diff --git a/src/hapsira/_math/optimize.py b/src/hapsira/core/_math/optimize.py similarity index 100% rename from src/hapsira/_math/optimize.py rename to src/hapsira/core/_math/optimize.py diff --git a/src/hapsira/_math/special.py b/src/hapsira/core/_math/special.py similarity index 100% rename from src/hapsira/_math/special.py rename to src/hapsira/core/_math/special.py diff --git a/src/hapsira/core/czml_utils.py b/src/hapsira/core/czml_utils.py index 54313646b..38597d1f0 100644 --- a/src/hapsira/core/czml_utils.py +++ b/src/hapsira/core/czml_utils.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from hapsira._math.linalg import norm +from ._math.linalg import norm @jit diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index a9be9bd08..18ea96ee7 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -8,10 +8,11 @@ import numpy as np from numpy import cos, cross, sin, sqrt -from hapsira._math.linalg import norm from hapsira.core.angles import E_to_nu, F_to_nu from hapsira.core.util import rotation_matrix +from ._math.linalg import norm + @jit def eccentricity_vector(k, r, v): diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index 921ba0c80..8e4348c5c 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -1,10 +1,11 @@ from numba import njit as jit import numpy as np -from hapsira._math.linalg import norm from hapsira.core.elements import coe_rotation_matrix, rv2coe from hapsira.core.util import planetocentric_to_AltAz +from ._math.linalg import norm + @jit def eclipse_function(k, u_, r_sec, R_sec, R_primary, umbra=True): diff --git a/src/hapsira/core/flybys.py b/src/hapsira/core/flybys.py index be8e8d656..c75ed565f 100644 --- a/src/hapsira/core/flybys.py +++ b/src/hapsira/core/flybys.py @@ -4,7 +4,7 @@ import numpy as np from numpy import cross -from hapsira._math.linalg import norm +from ._math.linalg import norm @jit diff --git a/src/hapsira/core/iod.py b/src/hapsira/core/iod.py index 601705ff2..2ae41e6ca 100644 --- a/src/hapsira/core/iod.py +++ b/src/hapsira/core/iod.py @@ -2,8 +2,8 @@ import numpy as np from numpy import cross, pi -from hapsira._math.linalg import norm -from hapsira._math.special import hyp2f1b, stumpff_c2 as c2, stumpff_c3 as c3 +from ._math.linalg import norm +from ._math.special import hyp2f1b, stumpff_c2 as c2, stumpff_c3 as c3 @jit diff --git a/src/hapsira/core/maneuver.py b/src/hapsira/core/maneuver.py index 548e5aea6..b081b6976 100644 --- a/src/hapsira/core/maneuver.py +++ b/src/hapsira/core/maneuver.py @@ -4,9 +4,10 @@ import numpy as np from numpy import cross -from hapsira._math.linalg import norm from hapsira.core.elements import coe_rotation_matrix, rv2coe, rv_pqw +from ._math.linalg import norm + @jit def hohmann(k, rv, r_f): diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index eb76a3582..5f6c28edc 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -1,9 +1,10 @@ from numba import njit as jit import numpy as np -from hapsira._math.linalg import norm from hapsira.core.events import line_of_sight as line_of_sight_fast +from ._math.linalg import norm + @jit def J2_perturbation(t0, state, k, J2, R): diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 876545b11..5dd3ecc00 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -1,8 +1,9 @@ import numpy as np -from hapsira._math.ivp import DOP853, solve_ivp from hapsira.core.propagation.base import func_twobody +from .._math.ivp import DOP853, solve_ivp + def cowell(k, r, v, tofs, rtol=1e-11, *, events=None, f=func_twobody): x, y, z = r diff --git a/src/hapsira/core/propagation/vallado.py b/src/hapsira/core/propagation/vallado.py index d52fc3a97..dde1e6e25 100644 --- a/src/hapsira/core/propagation/vallado.py +++ b/src/hapsira/core/propagation/vallado.py @@ -1,8 +1,8 @@ from numba import njit as jit import numpy as np -from hapsira._math.linalg import norm -from hapsira._math.special import stumpff_c2 as c2, stumpff_c3 as c3 +from .._math.linalg import norm +from .._math.special import stumpff_c2 as c2, stumpff_c3 as c3 @jit diff --git a/src/hapsira/core/spheroid_location.py b/src/hapsira/core/spheroid_location.py index dd6b54651..493ab4a65 100644 --- a/src/hapsira/core/spheroid_location.py +++ b/src/hapsira/core/spheroid_location.py @@ -3,7 +3,7 @@ from numba import njit as jit import numpy as np -from hapsira._math.linalg import norm +from ._math.linalg import norm @jit diff --git a/src/hapsira/core/thrust/change_a_inc.py b/src/hapsira/core/thrust/change_a_inc.py index 2dbb3abed..b59c8b76c 100644 --- a/src/hapsira/core/thrust/change_a_inc.py +++ b/src/hapsira/core/thrust/change_a_inc.py @@ -2,9 +2,10 @@ import numpy as np from numpy import cross -from hapsira._math.linalg import norm from hapsira.core.elements import circular_velocity +from .._math.linalg import norm + @jit def extra_quantities(k, a_0, a_f, inc_0, inc_f, f): diff --git a/src/hapsira/core/thrust/change_argp.py b/src/hapsira/core/thrust/change_argp.py index ffaf7d888..fb32e870e 100644 --- a/src/hapsira/core/thrust/change_argp.py +++ b/src/hapsira/core/thrust/change_argp.py @@ -2,9 +2,10 @@ import numpy as np from numpy import cross -from hapsira._math.linalg import norm from hapsira.core.elements import circular_velocity, rv2coe +from .._math.linalg import norm + @jit def delta_V(V, ecc, argp_0, argp_f, f, A): diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index dfa97c86b..b16c247c5 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -9,13 +9,14 @@ import numpy as np from numpy import cross -from hapsira._math.linalg import norm from hapsira.core.elements import ( circular_velocity, eccentricity_vector, rv2coe, ) +from .._math.linalg import norm + @jit def beta(ecc_0, ecc_f, inc_0, inc_f, argp): diff --git a/src/hapsira/earth/atmosphere/coesa62.py b/src/hapsira/earth/atmosphere/coesa62.py index ee6f5dcc4..7c165bae3 100644 --- a/src/hapsira/earth/atmosphere/coesa62.py +++ b/src/hapsira/earth/atmosphere/coesa62.py @@ -58,7 +58,7 @@ from astropy.utils.data import get_pkg_data_filename import numpy as np -from hapsira._math.integrate import quad +from hapsira.core._math.integrate import quad from hapsira.earth.atmosphere.base import COESA # Constants come from the original paper to achieve pure implementation diff --git a/src/hapsira/ephem.py b/src/hapsira/ephem.py index 5f46772ef..0c1895cf3 100644 --- a/src/hapsira/ephem.py +++ b/src/hapsira/ephem.py @@ -10,8 +10,8 @@ ) from astroquery.jplhorizons import Horizons -from hapsira._math.interpolate import interp1d, sinc_interp, spline_interp from hapsira.bodies import Earth +from hapsira.core._math.interpolate import interp1d, sinc_interp, spline_interp from hapsira.frames import Planes from hapsira.frames.util import get_frame from hapsira.twobody.sampling import EpochsArray diff --git a/src/hapsira/threebody/restricted.py b/src/hapsira/threebody/restricted.py index 76c5e1d10..57b42f7ec 100644 --- a/src/hapsira/threebody/restricted.py +++ b/src/hapsira/threebody/restricted.py @@ -7,7 +7,7 @@ from astropy import units as u import numpy as np -from hapsira._math.optimize import brentq +from hapsira.core._math.optimize import brentq from hapsira.util import norm diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index e0604b934..b8e720df8 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -4,7 +4,7 @@ from astropy.coordinates import get_body_barycentric_posvel import numpy as np -from hapsira._math.linalg import norm +from hapsira.core._math.linalg import norm from hapsira.core.events import ( eclipse_function as eclipse_function_fast, line_of_sight as line_of_sight_fast, diff --git a/src/hapsira/util.py b/src/hapsira/util.py index d1886867a..fb9c5add4 100644 --- a/src/hapsira/util.py +++ b/src/hapsira/util.py @@ -4,7 +4,7 @@ from astropy.time import Time import numpy as np -from hapsira._math.linalg import norm as norm_fast +from hapsira.core._math.linalg import norm as norm_fast from hapsira.core.util import alinspace as alinspace_fast diff --git a/tests/test_hyper.py b/tests/test_hyper.py index 1caf64fe7..c20471c5a 100644 --- a/tests/test_hyper.py +++ b/tests/test_hyper.py @@ -3,7 +3,7 @@ import pytest from scipy import special -from hapsira._math.special import hyp2f1b as hyp2f1 +from hapsira.core._math.special import hyp2f1b as hyp2f1 @pytest.mark.parametrize("x", np.linspace(0, 1, num=11)) diff --git a/tests/test_stumpff.py b/tests/test_stumpff.py index 4ee7b0be5..f20d7264d 100644 --- a/tests/test_stumpff.py +++ b/tests/test_stumpff.py @@ -1,7 +1,7 @@ from numpy import cos, cosh, sin, sinh from numpy.testing import assert_allclose -from hapsira._math.special import stumpff_c2 as c2, stumpff_c3 as c3 +from hapsira.core._math.special import stumpff_c2 as c2, stumpff_c3 as c3 def test_stumpff_functions_near_zero(): From 6357033fb06465092fd2bae6d0282ddbee9fc6f8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 27 Dec 2023 11:18:26 +0100 Subject: [PATCH 009/346] run slow tests too --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index e4f4feb27..40619fa4c 100644 --- a/Makefile +++ b/Makefile @@ -49,6 +49,6 @@ upload: done test: - DISPLAY= tox + DISPLAY= tox -e style,tests-fast,tests-slow,docs .PHONY: docs docker image release upload From a82ca68c78bf8bdcc0307e0386779d064969808b Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 27 Dec 2023 16:00:59 +0100 Subject: [PATCH 010/346] fix math import --- contrib/CR3BP/CR3BP.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/contrib/CR3BP/CR3BP.py b/contrib/CR3BP/CR3BP.py index 1b93cabc4..b50cfc03a 100644 --- a/contrib/CR3BP/CR3BP.py +++ b/contrib/CR3BP/CR3BP.py @@ -32,7 +32,7 @@ from numba import njit as jit import numpy as np -from hapsira._math.ivp import DOP853, solve_ivp +from hapsira.core._math.ivp import DOP853, solve_ivp @jit From 6c0fed6633503e697ff118827c8b94cecc518f7e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 27 Dec 2023 16:15:34 +0100 Subject: [PATCH 011/346] disable sort rules --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 441079dde..3389fb6c5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -131,6 +131,7 @@ select = [ ignore = [ "E501", # Line too long. Ignoring this so "black" manages line length. + "I001", # import sort ] exclude = [ @@ -143,7 +144,7 @@ exclude = [ [tool.ruff.isort] combine-as-imports = true -force-sort-within-sections = true +force-sort-within-sections = false known-first-party = ["hapsira"] [tool.ruff.mccabe] From 7d5b91b70c80850eddcf4f5621b6ffa032966e79 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 27 Dec 2023 16:15:47 +0100 Subject: [PATCH 012/346] more exports up --- src/hapsira/core/_jit.py | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/src/hapsira/core/_jit.py b/src/hapsira/core/_jit.py index 1b21e9d05..aed5794ca 100644 --- a/src/hapsira/core/_jit.py +++ b/src/hapsira/core/_jit.py @@ -1,10 +1,23 @@ from enum import Enum, auto -import os from typing import Callable +import os import numba as nb from numba import cuda -from poliastro.errors import JitError + +from hapsira.errors import JitError + + +__all__ = [ + "INLINE", + "NOPYTHON", + "PRECISIONS", + "TARGET", + "TARGETS", + "hjit", + "vjit", + "sjit", +] def _str2bool(value: str) -> bool: @@ -214,15 +227,3 @@ def wrapper(inner_func: Callable) -> Callable: return wrapper(outer_func) return wrapper - - -__all__ = [ - "INLINE", - "NOPYTHON", - "PRECISIONS", - "TARGET", - "TARGETS", - "hjit", - "vjit", - "sjit", -] From 13dfe5d8f4757d3753662a3aeab940b73db5bf1c Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 27 Dec 2023 16:16:13 +0100 Subject: [PATCH 013/346] jit becomes public --- src/hapsira/core/{_jit.py => jit.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename src/hapsira/core/{_jit.py => jit.py} (100%) diff --git a/src/hapsira/core/_jit.py b/src/hapsira/core/jit.py similarity index 100% rename from src/hapsira/core/_jit.py rename to src/hapsira/core/jit.py From 4555015431f33f00d90fd30743d62a1fd95477ca Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 27 Dec 2023 16:20:55 +0100 Subject: [PATCH 014/346] cleanup exports --- src/hapsira/core/_math/integrate.py | 4 +++- src/hapsira/core/_math/interpolate.py | 6 +++++- src/hapsira/core/_math/ivp.py | 5 ++++- src/hapsira/core/_math/linalg.py | 4 ++++ src/hapsira/core/_math/optimize.py | 4 +++- src/hapsira/core/_math/special.py | 6 ++++++ 6 files changed, 25 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/_math/integrate.py b/src/hapsira/core/_math/integrate.py index 49ad08a13..3a4565bad 100644 --- a/src/hapsira/core/_math/integrate.py +++ b/src/hapsira/core/_math/integrate.py @@ -1,3 +1,5 @@ from scipy.integrate import quad -__all__ = ["quad"] +__all__ = [ + "quad", +] diff --git a/src/hapsira/core/_math/interpolate.py b/src/hapsira/core/_math/interpolate.py index 711b1579b..b4614959a 100644 --- a/src/hapsira/core/_math/interpolate.py +++ b/src/hapsira/core/_math/interpolate.py @@ -1,7 +1,11 @@ import numpy as np from scipy.interpolate import interp1d -__all__ = ["interp1d", "spline_interp", "sinc_interp"] +__all__ = [ + "interp1d", + "spline_interp", + "sinc_interp", +] def spline_interp(y, x, u, *, kind="cubic"): diff --git a/src/hapsira/core/_math/ivp.py b/src/hapsira/core/_math/ivp.py index f5b9c01b1..29e4f7674 100644 --- a/src/hapsira/core/_math/ivp.py +++ b/src/hapsira/core/_math/ivp.py @@ -1,3 +1,6 @@ from scipy.integrate import DOP853, solve_ivp -__all__ = ["DOP853", "solve_ivp"] +__all__ = [ + "DOP853", + "solve_ivp", +] diff --git a/src/hapsira/core/_math/linalg.py b/src/hapsira/core/_math/linalg.py index a2845a43e..26b604270 100644 --- a/src/hapsira/core/_math/linalg.py +++ b/src/hapsira/core/_math/linalg.py @@ -1,6 +1,10 @@ from numba import njit as jit import numpy as np +__all__ = [ + "norm", +] + @jit def norm(arr): diff --git a/src/hapsira/core/_math/optimize.py b/src/hapsira/core/_math/optimize.py index 9ddcedca7..4033e6080 100644 --- a/src/hapsira/core/_math/optimize.py +++ b/src/hapsira/core/_math/optimize.py @@ -1,3 +1,5 @@ from scipy.optimize import brentq -__all__ = ["brentq"] +__all__ = [ + "brentq", +] diff --git a/src/hapsira/core/_math/special.py b/src/hapsira/core/_math/special.py index 6f704a965..375cce97f 100644 --- a/src/hapsira/core/_math/special.py +++ b/src/hapsira/core/_math/special.py @@ -3,6 +3,12 @@ from numba import njit as jit import numpy as np +__all__ = [ + "hyp2f1b", + "stumpff_c2", + "stumpff_c3", +] + @jit def hyp2f1b(x): From 2b6f3e6ca7381aad8b15076b843d25b932c2d9bf Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 28 Dec 2023 11:06:25 +0100 Subject: [PATCH 015/346] rename name module --- src/hapsira/core/{_math => math}/__init__.py | 0 src/hapsira/core/{_math => math}/integrate.py | 0 src/hapsira/core/{_math => math}/interpolate.py | 0 src/hapsira/core/{_math => math}/ivp.py | 0 src/hapsira/core/{_math => math}/linalg.py | 0 src/hapsira/core/{_math => math}/optimize.py | 0 src/hapsira/core/{_math => math}/special.py | 0 7 files changed, 0 insertions(+), 0 deletions(-) rename src/hapsira/core/{_math => math}/__init__.py (100%) rename src/hapsira/core/{_math => math}/integrate.py (100%) rename src/hapsira/core/{_math => math}/interpolate.py (100%) rename src/hapsira/core/{_math => math}/ivp.py (100%) rename src/hapsira/core/{_math => math}/linalg.py (100%) rename src/hapsira/core/{_math => math}/optimize.py (100%) rename src/hapsira/core/{_math => math}/special.py (100%) diff --git a/src/hapsira/core/_math/__init__.py b/src/hapsira/core/math/__init__.py similarity index 100% rename from src/hapsira/core/_math/__init__.py rename to src/hapsira/core/math/__init__.py diff --git a/src/hapsira/core/_math/integrate.py b/src/hapsira/core/math/integrate.py similarity index 100% rename from src/hapsira/core/_math/integrate.py rename to src/hapsira/core/math/integrate.py diff --git a/src/hapsira/core/_math/interpolate.py b/src/hapsira/core/math/interpolate.py similarity index 100% rename from src/hapsira/core/_math/interpolate.py rename to src/hapsira/core/math/interpolate.py diff --git a/src/hapsira/core/_math/ivp.py b/src/hapsira/core/math/ivp.py similarity index 100% rename from src/hapsira/core/_math/ivp.py rename to src/hapsira/core/math/ivp.py diff --git a/src/hapsira/core/_math/linalg.py b/src/hapsira/core/math/linalg.py similarity index 100% rename from src/hapsira/core/_math/linalg.py rename to src/hapsira/core/math/linalg.py diff --git a/src/hapsira/core/_math/optimize.py b/src/hapsira/core/math/optimize.py similarity index 100% rename from src/hapsira/core/_math/optimize.py rename to src/hapsira/core/math/optimize.py diff --git a/src/hapsira/core/_math/special.py b/src/hapsira/core/math/special.py similarity index 100% rename from src/hapsira/core/_math/special.py rename to src/hapsira/core/math/special.py From 8488ca9036c2259bb6d42076f01d0ad3f487068e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 28 Dec 2023 11:17:53 +0100 Subject: [PATCH 016/346] fixed math imports --- contrib/CR3BP/CR3BP.py | 2 +- src/hapsira/core/czml_utils.py | 2 +- src/hapsira/core/elements.py | 2 +- src/hapsira/core/events.py | 2 +- src/hapsira/core/flybys.py | 2 +- src/hapsira/core/iod.py | 4 ++-- src/hapsira/core/maneuver.py | 2 +- src/hapsira/core/perturbations.py | 2 +- src/hapsira/core/propagation/cowell.py | 2 +- src/hapsira/core/propagation/vallado.py | 4 ++-- src/hapsira/core/spheroid_location.py | 2 +- src/hapsira/core/thrust/change_a_inc.py | 2 +- src/hapsira/core/thrust/change_argp.py | 2 +- src/hapsira/core/thrust/change_ecc_inc.py | 2 +- src/hapsira/earth/atmosphere/coesa62.py | 2 +- src/hapsira/ephem.py | 2 +- src/hapsira/threebody/restricted.py | 2 +- src/hapsira/twobody/events.py | 2 +- src/hapsira/util.py | 2 +- tests/test_hyper.py | 2 +- tests/test_stumpff.py | 2 +- 21 files changed, 23 insertions(+), 23 deletions(-) diff --git a/contrib/CR3BP/CR3BP.py b/contrib/CR3BP/CR3BP.py index b50cfc03a..237f8e62f 100644 --- a/contrib/CR3BP/CR3BP.py +++ b/contrib/CR3BP/CR3BP.py @@ -32,7 +32,7 @@ from numba import njit as jit import numpy as np -from hapsira.core._math.ivp import DOP853, solve_ivp +from hapsira.core.math.ivp import DOP853, solve_ivp @jit diff --git a/src/hapsira/core/czml_utils.py b/src/hapsira/core/czml_utils.py index 38597d1f0..d479c3bed 100644 --- a/src/hapsira/core/czml_utils.py +++ b/src/hapsira/core/czml_utils.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from ._math.linalg import norm +from .math.linalg import norm @jit diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 18ea96ee7..67f12d5e3 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -11,7 +11,7 @@ from hapsira.core.angles import E_to_nu, F_to_nu from hapsira.core.util import rotation_matrix -from ._math.linalg import norm +from .math.linalg import norm @jit diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index 8e4348c5c..44b736b63 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -4,7 +4,7 @@ from hapsira.core.elements import coe_rotation_matrix, rv2coe from hapsira.core.util import planetocentric_to_AltAz -from ._math.linalg import norm +from .math.linalg import norm @jit diff --git a/src/hapsira/core/flybys.py b/src/hapsira/core/flybys.py index c75ed565f..4849a3d44 100644 --- a/src/hapsira/core/flybys.py +++ b/src/hapsira/core/flybys.py @@ -4,7 +4,7 @@ import numpy as np from numpy import cross -from ._math.linalg import norm +from .math.linalg import norm @jit diff --git a/src/hapsira/core/iod.py b/src/hapsira/core/iod.py index 2ae41e6ca..a6ed6f7e3 100644 --- a/src/hapsira/core/iod.py +++ b/src/hapsira/core/iod.py @@ -2,8 +2,8 @@ import numpy as np from numpy import cross, pi -from ._math.linalg import norm -from ._math.special import hyp2f1b, stumpff_c2 as c2, stumpff_c3 as c3 +from .math.linalg import norm +from .math.special import hyp2f1b, stumpff_c2 as c2, stumpff_c3 as c3 @jit diff --git a/src/hapsira/core/maneuver.py b/src/hapsira/core/maneuver.py index b081b6976..b7f7936f0 100644 --- a/src/hapsira/core/maneuver.py +++ b/src/hapsira/core/maneuver.py @@ -6,7 +6,7 @@ from hapsira.core.elements import coe_rotation_matrix, rv2coe, rv_pqw -from ._math.linalg import norm +from .math.linalg import norm @jit diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index 5f6c28edc..3b1800d73 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -3,7 +3,7 @@ from hapsira.core.events import line_of_sight as line_of_sight_fast -from ._math.linalg import norm +from .math.linalg import norm @jit diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 5dd3ecc00..8c8bae666 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -2,7 +2,7 @@ from hapsira.core.propagation.base import func_twobody -from .._math.ivp import DOP853, solve_ivp +from ..math.ivp import DOP853, solve_ivp def cowell(k, r, v, tofs, rtol=1e-11, *, events=None, f=func_twobody): diff --git a/src/hapsira/core/propagation/vallado.py b/src/hapsira/core/propagation/vallado.py index dde1e6e25..c6b6f2914 100644 --- a/src/hapsira/core/propagation/vallado.py +++ b/src/hapsira/core/propagation/vallado.py @@ -1,8 +1,8 @@ from numba import njit as jit import numpy as np -from .._math.linalg import norm -from .._math.special import stumpff_c2 as c2, stumpff_c3 as c3 +from ..math.linalg import norm +from ..math.special import stumpff_c2 as c2, stumpff_c3 as c3 @jit diff --git a/src/hapsira/core/spheroid_location.py b/src/hapsira/core/spheroid_location.py index 493ab4a65..3f49bd38e 100644 --- a/src/hapsira/core/spheroid_location.py +++ b/src/hapsira/core/spheroid_location.py @@ -3,7 +3,7 @@ from numba import njit as jit import numpy as np -from ._math.linalg import norm +from .math.linalg import norm @jit diff --git a/src/hapsira/core/thrust/change_a_inc.py b/src/hapsira/core/thrust/change_a_inc.py index b59c8b76c..a97c66fb2 100644 --- a/src/hapsira/core/thrust/change_a_inc.py +++ b/src/hapsira/core/thrust/change_a_inc.py @@ -4,7 +4,7 @@ from hapsira.core.elements import circular_velocity -from .._math.linalg import norm +from ..math.linalg import norm @jit diff --git a/src/hapsira/core/thrust/change_argp.py b/src/hapsira/core/thrust/change_argp.py index fb32e870e..4fc8da570 100644 --- a/src/hapsira/core/thrust/change_argp.py +++ b/src/hapsira/core/thrust/change_argp.py @@ -4,7 +4,7 @@ from hapsira.core.elements import circular_velocity, rv2coe -from .._math.linalg import norm +from ..math.linalg import norm @jit diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index b16c247c5..61fa04ea1 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -15,7 +15,7 @@ rv2coe, ) -from .._math.linalg import norm +from ..math.linalg import norm @jit diff --git a/src/hapsira/earth/atmosphere/coesa62.py b/src/hapsira/earth/atmosphere/coesa62.py index 7c165bae3..225772ea0 100644 --- a/src/hapsira/earth/atmosphere/coesa62.py +++ b/src/hapsira/earth/atmosphere/coesa62.py @@ -58,7 +58,7 @@ from astropy.utils.data import get_pkg_data_filename import numpy as np -from hapsira.core._math.integrate import quad +from hapsira.core.math.integrate import quad from hapsira.earth.atmosphere.base import COESA # Constants come from the original paper to achieve pure implementation diff --git a/src/hapsira/ephem.py b/src/hapsira/ephem.py index 0c1895cf3..853583028 100644 --- a/src/hapsira/ephem.py +++ b/src/hapsira/ephem.py @@ -11,7 +11,7 @@ from astroquery.jplhorizons import Horizons from hapsira.bodies import Earth -from hapsira.core._math.interpolate import interp1d, sinc_interp, spline_interp +from hapsira.core.math.interpolate import interp1d, sinc_interp, spline_interp from hapsira.frames import Planes from hapsira.frames.util import get_frame from hapsira.twobody.sampling import EpochsArray diff --git a/src/hapsira/threebody/restricted.py b/src/hapsira/threebody/restricted.py index 57b42f7ec..58b37ba67 100644 --- a/src/hapsira/threebody/restricted.py +++ b/src/hapsira/threebody/restricted.py @@ -7,7 +7,7 @@ from astropy import units as u import numpy as np -from hapsira.core._math.optimize import brentq +from hapsira.core.math.optimize import brentq from hapsira.util import norm diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index b8e720df8..ab54e952c 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -4,7 +4,7 @@ from astropy.coordinates import get_body_barycentric_posvel import numpy as np -from hapsira.core._math.linalg import norm +from hapsira.core.math.linalg import norm from hapsira.core.events import ( eclipse_function as eclipse_function_fast, line_of_sight as line_of_sight_fast, diff --git a/src/hapsira/util.py b/src/hapsira/util.py index fb9c5add4..5b56137d0 100644 --- a/src/hapsira/util.py +++ b/src/hapsira/util.py @@ -4,7 +4,7 @@ from astropy.time import Time import numpy as np -from hapsira.core._math.linalg import norm as norm_fast +from hapsira.core.math.linalg import norm as norm_fast from hapsira.core.util import alinspace as alinspace_fast diff --git a/tests/test_hyper.py b/tests/test_hyper.py index c20471c5a..edeea790a 100644 --- a/tests/test_hyper.py +++ b/tests/test_hyper.py @@ -3,7 +3,7 @@ import pytest from scipy import special -from hapsira.core._math.special import hyp2f1b as hyp2f1 +from hapsira.core.math.special import hyp2f1b as hyp2f1 @pytest.mark.parametrize("x", np.linspace(0, 1, num=11)) diff --git a/tests/test_stumpff.py b/tests/test_stumpff.py index f20d7264d..60ccd6830 100644 --- a/tests/test_stumpff.py +++ b/tests/test_stumpff.py @@ -1,7 +1,7 @@ from numpy import cos, cosh, sin, sinh from numpy.testing import assert_allclose -from hapsira.core._math.special import stumpff_c2 as c2, stumpff_c3 as c3 +from hapsira.core.math.special import stumpff_c2 as c2, stumpff_c3 as c3 def test_stumpff_functions_near_zero(): From 38ffa5a988e29b3b4c5a88383d005aef0cd89ab5 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 28 Dec 2023 20:07:05 +0100 Subject: [PATCH 017/346] debug module --- src/hapsira/debug.py | 37 +++++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) create mode 100644 src/hapsira/debug.py diff --git a/src/hapsira/debug.py b/src/hapsira/debug.py new file mode 100644 index 000000000..b348ab00b --- /dev/null +++ b/src/hapsira/debug.py @@ -0,0 +1,37 @@ +import logging +import os + +__all__ = [ + "DEBUG", + "LOGLEVEL", + "get_environ_switch", + "logger", +] + + +def get_environ_switch(name: str, default: bool) -> bool: + """ + Helper for parsing environment variables + """ + + value = os.environ.get(name, "1" if default else "0") + + if value.strip().lower() in ("true", "1", "yes"): + return True + if value.strip().lower() in ("false", "0", "no"): + return False + + raise ValueError(f'can not convert value "{value:s}" to bool') + + +DEBUG = get_environ_switch("HAPSIRA_DEBUG", default=False) +LOGLEVEL = os.environ.get("HAPSIRA_LOGLEVEL", "WARNING") +if LOGLEVEL not in ("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL", "NOTSET"): + raise ValueError(f'Unknown loglevel "{LOGLEVEL:s}"') + +logger = logging.getLogger("hapsira") +if LOGLEVEL != "NOTSET": + logger.setLevel(logging.DEBUG if DEBUG else getattr(logging, LOGLEVEL)) + +logger.debug("debug mode: %s", "on" if DEBUG else "off") +logger.debug("logging level: %s", LOGLEVEL) From db5e885ae801060841f8db375b5b089db1e745eb Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 28 Dec 2023 20:07:59 +0100 Subject: [PATCH 018/346] added logging to jit wrappers; moved env parsing to log module; fix exception raise missing --- src/hapsira/core/jit.py | 61 +++++++++++++++++++++++++++-------------- 1 file changed, 41 insertions(+), 20 deletions(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index aed5794ca..629d32f6e 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -5,6 +5,7 @@ import numba as nb from numba import cuda +from hapsira.debug import get_environ_switch, logger from hapsira.errors import JitError @@ -20,18 +21,6 @@ ] -def _str2bool(value: str) -> bool: - """ - Helper for parsing environment variables - """ - - if value.strip().lower() in ("true", "1", "yes"): - return True - if value.strip().lower() in ("false", "0", "no"): - return False - raise ValueError(f'can not convert value "{value:s}" to bool') - - class TARGETS(Enum): """ JIT targets @@ -75,14 +64,17 @@ def get_current(cls): TARGET = TARGETS.get_current() +logger.debug("jit option target: %s", TARGET.name) -INLINE = _str2bool( - os.environ.get("HAPSIRA_INLINE", "1" if TARGET is TARGET.cuda else "0") +INLINE = get_environ_switch( + "HAPSIRA_INLINE", default=TARGET is TARGET.cuda ) # currently only relevant for helpers on cpu and parallel targets +logger.debug("jit option inline: %s", "yes" if INLINE else "no") -NOPYTHON = _str2bool( - os.environ.get("HAPSIRA_NOPYTHON", "1") +NOPYTHON = get_environ_switch( + "HAPSIRA_NOPYTHON", default=True ) # only for debugging, True by default +logger.debug("jit option nopython: %s", "yes" if NOPYTHON else "no") PRECISIONS = ("f4", "f8") # TODO allow f2, i.e. half, for CUDA at least? @@ -95,11 +87,15 @@ def _parse_signatures(signature: str) -> str | list[str]: if not any( notation in signature for notation in ("f", "V") ): # leave this signature as it is - # TODO warn! + logger.warning( + "jit signature: no special notation, not parsing (%s)", signature + ) return signature if any(level in signature for level in PRECISIONS): # leave this signature as it is - # TODO warn! + logger.warning( + "jit signature: precision specified, not parsing (%s)", signature + ) return signature signature = signature.replace( @@ -141,11 +137,18 @@ def wrapper(inner_func: Callable) -> Callable: inline=INLINE, ) else: - JitError( + raise JitError( f'unknown target "{repr(TARGET):s}"; known targets are {repr(TARGETS):s}' ) cfg.update(kwargs) + logger.debug( + "hjit: %s, %s, %s", + getattr(inner_func, "__name__", repr(inner_func)), + repr(args), + repr(cfg), + ) + return wjit( *args, **cfg, @@ -182,8 +185,19 @@ def wrapper(inner_func: Callable) -> Callable: ) if TARGET is not TARGETS.cuda: cfg["nopython"] = NOPYTHON + elif TARGET not in TARGETS: + raise JitError( + f'unknown target "{repr(TARGET):s}"; known targets are {repr(TARGETS):s}' + ) cfg.update(kwargs) + logger.debug( + "vjit: %s, %s, %s", + getattr(inner_func, "__name__", repr(inner_func)), + repr(args), + repr(cfg), + ) + return nb.vectorize( *args, **cfg, @@ -215,8 +229,15 @@ def wrapper(inner_func: Callable) -> Callable: cfg = dict( nopython=NOPYTHON, inline="always" if INLINE else "never", + **kwargs, + ) + + logger.debug( + "sjit: %s, %s, %s", + getattr(inner_func, "__name__", repr(inner_func)), + repr(args), + repr(cfg), ) - cfg.update(kwargs) return nb.jit( *args, From b8cd85714b0ea90b5e0d832662d3ff584a0c4dcf Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 29 Dec 2023 00:22:27 +0100 Subject: [PATCH 019/346] new settings module --- src/hapsira/settings.py | 154 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 src/hapsira/settings.py diff --git a/src/hapsira/settings.py b/src/hapsira/settings.py new file mode 100644 index 000000000..d24505391 --- /dev/null +++ b/src/hapsira/settings.py @@ -0,0 +1,154 @@ +import os +from typing import Any, Generator, Optional, Type + +__all__ = [ + "Setting", + "Settings", + "settings", +] + + +def _str2bool(value: str) -> bool: + """ + Helper for parsing environment variables + """ + + if value.strip().lower() in ("true", "1", "yes", "y"): + return True + if value.strip().lower() in ("false", "0", "no", "n"): + return False + + raise ValueError(f'can not convert value "{value:s}" to bool') + + +class Setting: + """ + Holds one setting settable by user before sub-module import + """ + + def __init__(self, name: str, default: Any, options: Optional[tuple[Any]] = None): + self._name = name + self._type = type(default) + self._value = default + self._options = options + self._check_env() + + def _check_env(self): + """ + Check for environment variables + """ + value = os.environ.get(f"HAPSIRA_{self._name:s}") + if value is None: + return + if self._type is bool: + value = _str2bool(value) + self.value = value # Run through setter for checks! + + @property + def name(self) -> str: + """ + Return name of setting + """ + return self._name + + @property + def type_(self) -> Type: + """ + Return type of setting + """ + return self._type + + @property + def options(self) -> Optional[tuple[Any]]: + """ + Return options for value + """ + return self._options + + @property + def value(self) -> Any: + """ + Change value of setting + """ + return self._value + + @value.setter + def value(self, new_value: Any): + """ + Return value of setting + """ + if not isinstance(new_value, self._type): + raise TypeError( + f'"{repr():s}" has type "{repr():s}", expected type"{repr():s}"' + ) + if self._options is not None and new_value not in self._options: + raise ValueError( + f'value "{repr(new_value):s}" not a valid option, valid options are "{repr(self._options):s}"' + ) + self._value = new_value + + +class Settings: + """ + Holds settings settable by user before sub-module import + """ + + def __init__(self): + self._settings = {} + self._add( + Setting( + "DEBUG", + False, + ) + ) + self._add( + Setting( + "LOGLEVEL", + "NOTSET", + options=("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL", "NOTSET"), + ) + ) + self._add( + Setting( + "TARGET", + "cpu", + options=("cpu", "parallel", "cuda"), + ) + ) + self._add( + Setting( + "INLINE", + self["TARGET"].value == "cuda", + ) + ) + self._add( + Setting( + "NOPYTHON", + True, + ) + ) + + def _add(self, setting: Setting): + """ + Add new setting + """ + self._settings[setting.name] = setting + + def __getitem__(self, name: str) -> Setting: + """ + Return setting by name + """ + if name not in self._settings.keys(): + raise KeyError( + f'setting "{name:s}" unknown, possible settings are {repr(list(self._settings.keys())):s}' + ) + return self._settings[name] + + def keys(self) -> Generator: + """ + Generator of all setting names + """ + return (name for name in self._settings.keys()) + + +settings = Settings() From 14010977b29119f2ce5c59f25453e4a86c7da14e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 29 Dec 2023 00:22:49 +0100 Subject: [PATCH 020/346] rely on new settings module --- src/hapsira/debug.py | 40 +++++++++++----------------------------- 1 file changed, 11 insertions(+), 29 deletions(-) diff --git a/src/hapsira/debug.py b/src/hapsira/debug.py index b348ab00b..de206d539 100644 --- a/src/hapsira/debug.py +++ b/src/hapsira/debug.py @@ -1,37 +1,19 @@ import logging -import os + +from hapsira.settings import settings __all__ = [ - "DEBUG", - "LOGLEVEL", - "get_environ_switch", "logger", ] - -def get_environ_switch(name: str, default: bool) -> bool: - """ - Helper for parsing environment variables - """ - - value = os.environ.get(name, "1" if default else "0") - - if value.strip().lower() in ("true", "1", "yes"): - return True - if value.strip().lower() in ("false", "0", "no"): - return False - - raise ValueError(f'can not convert value "{value:s}" to bool') - - -DEBUG = get_environ_switch("HAPSIRA_DEBUG", default=False) -LOGLEVEL = os.environ.get("HAPSIRA_LOGLEVEL", "WARNING") -if LOGLEVEL not in ("DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL", "NOTSET"): - raise ValueError(f'Unknown loglevel "{LOGLEVEL:s}"') - logger = logging.getLogger("hapsira") -if LOGLEVEL != "NOTSET": - logger.setLevel(logging.DEBUG if DEBUG else getattr(logging, LOGLEVEL)) -logger.debug("debug mode: %s", "on" if DEBUG else "off") -logger.debug("logging level: %s", LOGLEVEL) +if settings["LOGLEVEL"].value != "NOTSET": + logger.setLevel( + logging.DEBUG + if settings["DEBUG"].value + else getattr(logging, settings["LOGLEVEL"].value) + ) + +logger.debug("logging level: %s", logging.getLevelName(logger.level)) +logger.debug("debug mode: %s", "on" if settings["DEBUG"].value else "off") From a34efe05a1ed8d323661dffa9d52440e370d7dbe Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 29 Dec 2023 00:23:49 +0100 Subject: [PATCH 021/346] rely on new settings module --- src/hapsira/core/jit.py | 103 ++++++++-------------------------------- 1 file changed, 21 insertions(+), 82 deletions(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 629d32f6e..00897f205 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -1,80 +1,27 @@ -from enum import Enum, auto from typing import Callable -import os import numba as nb from numba import cuda -from hapsira.debug import get_environ_switch, logger +from hapsira.debug import logger from hapsira.errors import JitError +from hapsira.settings import settings __all__ = [ - "INLINE", - "NOPYTHON", "PRECISIONS", - "TARGET", - "TARGETS", "hjit", "vjit", "sjit", ] -class TARGETS(Enum): - """ - JIT targets - """ - - cpu = auto() - parallel = auto() - cuda = auto() - # numba 0.54.0, 19 August 2021, removed AMD ROCm target - - @classmethod - def get_default(cls): - """ - Default JIT target - """ - - return cls.cpu - - @classmethod - def get_current(cls): - """ - Current JIT target - """ - - name = os.environ.get("HAPSIRA_TARGET", None) - - if name is None: - target = cls.get_default() - else: - try: - target = cls[name] - except KeyError as e: - raise JitError( - f'unknown target "{name:s}"; known targets are {repr(cls):s}' - ) from e - - if target is cls.cuda and not cuda.is_available(): - raise JitError('selected target "cuda" is not available') - - return target - - -TARGET = TARGETS.get_current() -logger.debug("jit option target: %s", TARGET.name) - -INLINE = get_environ_switch( - "HAPSIRA_INLINE", default=TARGET is TARGET.cuda -) # currently only relevant for helpers on cpu and parallel targets -logger.debug("jit option inline: %s", "yes" if INLINE else "no") +logger.debug("jit target: %s", settings["TARGET"].value) +if settings["TARGET"].value == "cuda" and not cuda.is_available(): + raise JitError('selected target "cuda" is not available') -NOPYTHON = get_environ_switch( - "HAPSIRA_NOPYTHON", default=True -) # only for debugging, True by default -logger.debug("jit option nopython: %s", "yes" if NOPYTHON else "no") +logger.debug("jit inline: %s", "yes" if settings["INLINE"].value else "no") +logger.debug("jit nopython: %s", "yes" if settings["NOPYTHON"].value else "no") PRECISIONS = ("f4", "f8") # TODO allow f2, i.e. half, for CUDA at least? @@ -124,26 +71,22 @@ def wrapper(inner_func: Callable) -> Callable: Applies JIT """ - if TARGET in (TARGETS.cpu, TARGETS.parallel): - wjit = nb.jit - cfg = dict( - nopython=NOPYTHON, - inline="always" if INLINE else "never", - ) - elif TARGET is TARGETS.cuda: + if settings["TARGET"].value == "cuda": wjit = cuda.jit cfg = dict( device=True, - inline=INLINE, + inline=settings["INLINE"].value, ) else: - raise JitError( - f'unknown target "{repr(TARGET):s}"; known targets are {repr(TARGETS):s}' + wjit = nb.jit + cfg = dict( + nopython=settings["NOPYTHON"].value, + inline="always" if settings["INLINE"].value else "never", ) cfg.update(kwargs) logger.debug( - "hjit: %s, %s, %s", + "hjit: func=%s, args=%s, kwargs=%s", getattr(inner_func, "__name__", repr(inner_func)), repr(args), repr(cfg), @@ -181,18 +124,14 @@ def wrapper(inner_func: Callable) -> Callable: """ cfg = dict( - target=TARGET.name, + target=settings["TARGET"].value, ) - if TARGET is not TARGETS.cuda: - cfg["nopython"] = NOPYTHON - elif TARGET not in TARGETS: - raise JitError( - f'unknown target "{repr(TARGET):s}"; known targets are {repr(TARGETS):s}' - ) + if settings["TARGET"].value != "cuda": + cfg["nopython"] = settings["NOPYTHON"].value cfg.update(kwargs) logger.debug( - "vjit: %s, %s, %s", + "vjit: func=%s, args=%s, kwargs=%s", getattr(inner_func, "__name__", repr(inner_func)), repr(args), repr(cfg), @@ -227,13 +166,13 @@ def wrapper(inner_func: Callable) -> Callable: """ cfg = dict( - nopython=NOPYTHON, - inline="always" if INLINE else "never", + nopython=settings["NOPYTHON"].value, + inline="always" if settings["INLINE"].value else "never", **kwargs, ) logger.debug( - "sjit: %s, %s, %s", + "sjit: func=%s, args=%s, kwargs=%s", getattr(inner_func, "__name__", repr(inner_func)), repr(args), repr(cfg), From 2220f34cbae93c53ec3c3886a73dd8308cdaf7fd Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 29 Dec 2023 00:42:07 +0100 Subject: [PATCH 022/346] shortcut for changing setting values --- src/hapsira/settings.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/src/hapsira/settings.py b/src/hapsira/settings.py index d24505391..e200350e1 100644 --- a/src/hapsira/settings.py +++ b/src/hapsira/settings.py @@ -134,16 +134,30 @@ def _add(self, setting: Setting): """ self._settings[setting.name] = setting + def _validate(self, name: str): + """ + Validate name + """ + if name in self._settings.keys(): + return + raise KeyError( + f'setting "{name:s}" unknown, possible settings are {repr(list(self._settings.keys())):s}' + ) + def __getitem__(self, name: str) -> Setting: """ Return setting by name """ - if name not in self._settings.keys(): - raise KeyError( - f'setting "{name:s}" unknown, possible settings are {repr(list(self._settings.keys())):s}' - ) + self._validate(name) return self._settings[name] + def __setitem__(self, name: str, new_value: Any): + """ + Return setting by name + """ + self._validate(name) + self._settings[name].value = new_value + def keys(self) -> Generator: """ Generator of all setting names From 11ff8dfaa440b21ef7bec8879ded466d0f08cbfc Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 29 Dec 2023 00:47:39 +0100 Subject: [PATCH 023/346] fix repr --- src/hapsira/settings.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hapsira/settings.py b/src/hapsira/settings.py index e200350e1..18a1b85dd 100644 --- a/src/hapsira/settings.py +++ b/src/hapsira/settings.py @@ -79,11 +79,11 @@ def value(self, new_value: Any): """ if not isinstance(new_value, self._type): raise TypeError( - f'"{repr():s}" has type "{repr():s}", expected type"{repr():s}"' + f"{repr(new_value):s} has type {repr(type(new_value)):s}, expected type {repr(self._type):s}" ) if self._options is not None and new_value not in self._options: raise ValueError( - f'value "{repr(new_value):s}" not a valid option, valid options are "{repr(self._options):s}"' + f"value {repr(new_value):s} not a valid option, valid options are {repr(self._options):s}" ) self._value = new_value From 3baa744150f515b705efdef540f39e433aa2d03a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 29 Dec 2023 14:20:21 +0100 Subject: [PATCH 024/346] expand signatures in sjit --- src/hapsira/core/jit.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 00897f205..ef4568115 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -160,6 +160,9 @@ def sjit(*args, **kwargs) -> Callable: else: pass + if len(args) > 0 and isinstance(args[0], str): + args = _parse_signatures(args[0]), *args[1:] + def wrapper(inner_func: Callable) -> Callable: """ Applies JIT From d3fd2b343f1bac79694b32eac93288c1247aa543 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 29 Dec 2023 14:31:47 +0100 Subject: [PATCH 025/346] fix scoping --- src/hapsira/core/jit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index ef4568115..e130547c0 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -158,7 +158,7 @@ def sjit(*args, **kwargs) -> Callable: outer_func = args[0] args = tuple() else: - pass + outer_func = None if len(args) > 0 and isinstance(args[0], str): args = _parse_signatures(args[0]), *args[1:] From 382e667640de633ffd0cc6d37e4a0863b78456a9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 12:08:01 +0100 Subject: [PATCH 026/346] core draft --- src/hapsira/core/angles.py | 395 +++++++++++++++++++++++++------------ 1 file changed, 266 insertions(+), 129 deletions(-) diff --git a/src/hapsira/core/angles.py b/src/hapsira/core/angles.py index 5e139f99f..d9f0b921d 100644 --- a/src/hapsira/core/angles.py +++ b/src/hapsira/core/angles.py @@ -1,53 +1,159 @@ -from numba import njit as jit -import numpy as np +from math import ( + asinh, + atan, + atan2, + atanh, + cos, + cosh, + nan, + pi, + sin, + sinh, + sqrt, + tan, + tanh, +) + +from .jit import hjit, vjit + +_TOL = 1.48e-08 + + +@hjit("f(f,f)") +def E_to_M_hf(E, ecc): + r"""Mean anomaly from eccentric anomaly. -@jit -def _kepler_equation(E, M, ecc): - return E_to_M(E, ecc) - M + .. versionadded:: 0.4.0 + Parameters + ---------- + E : float + Eccentric anomaly in radians. + ecc : float + Eccentricity. + + Returns + ------- + M : float + Mean anomaly. + + Warnings + -------- + The mean anomaly will be outside of (-π, π] + if the eccentric anomaly is. + No validation or wrapping is performed. -@jit -def _kepler_equation_prime(E, M, ecc): - return 1 - ecc * np.cos(E) + Notes + ----- + The implementation uses the plain original Kepler equation: + .. math:: + M = E - e \sin{E} -@jit -def _kepler_equation_hyper(F, M, ecc): - return F_to_M(F, ecc) - M + """ + M = E - ecc * sin(E) + return M -@jit -def _kepler_equation_prime_hyper(F, M, ecc): - return ecc * np.cosh(F) - 1 +@vjit("f(f,f)") +def E_to_M_vf(E, ecc): + """ + Vectorized E_to_M + """ + return E_to_M_hf(E, ecc) -def newton_factory(func, fprime): - @jit - def jit_newton_wrapper(x0, args=(), tol=1.48e-08, maxiter=50): - p0 = float(x0) - for _ in range(maxiter): - fval = func(p0, *args) - fder = fprime(p0, *args) - newton_step = fval / fder - p = p0 - newton_step - if abs(p - p0) < tol: - return p - p0 = p - return np.nan +@hjit("f(f,f)") +def F_to_M_hf(F, ecc): + r"""Mean anomaly from hyperbolic anomaly. - return jit_newton_wrapper + Parameters + ---------- + F : float + Hyperbolic anomaly. + ecc : float + Eccentricity (>1). + Returns + ------- + M : float + Mean anomaly. -_newton_elliptic = newton_factory(_kepler_equation, _kepler_equation_prime) -_newton_hyperbolic = newton_factory( - _kepler_equation_hyper, _kepler_equation_prime_hyper -) + Notes + ----- + As noted in [5]_, by manipulating + the parametric equations of the hyperbola + we can derive a quantity that is equivalent + to the mean anomaly in the elliptic case: + + .. math:: + + M = e \sinh{F} - F + + """ + M = ecc * sinh(F) - F + return M + + +@vjit("f(f,f)") +def F_to_M_vf(F, ecc): + """ + Vectorized F_to_M + """ + + return F_to_M_hf(F, ecc) + + +@hjit("f(f,f,f)") +def _kepler_equation_hf(E, M, ecc): + return E_to_M_hf(E, ecc) - M + + +@hjit("f(f,f,f)") +def _kepler_equation_prime_hf(E, M, ecc): + return 1 - ecc * cos(E) + + +@hjit("f(f,f,f)") +def _kepler_equation_hyper_hf(F, M, ecc): + return F_to_M_hf(F, ecc) - M -@jit -def D_to_nu(D): +@hjit("f(f,f,f)") +def _kepler_equation_prime_hyper_hf(F, M, ecc): + return ecc * cosh(F) - 1 + + +@hjit("f(f,f,f,f,i64)") +def _newton_elliptic_hf(p0, M, ecc, tol, maxiter): + for _ in range(maxiter): + fval = _kepler_equation_hf(p0, M, ecc) + fder = _kepler_equation_prime_hf(p0, M, ecc) + newton_step = fval / fder + p = p0 - newton_step + if abs(p - p0) < tol: + return p + p0 = p + return nan + + +@hjit("f(f,f,f,f,i64)") +def _newton_hyperbolic_hf(p0, M, ecc, tol, maxiter): + for _ in range(maxiter): + fval = _kepler_equation_hyper_hf(p0, M, ecc) + fder = _kepler_equation_prime_hyper_hf(p0, M, ecc) + newton_step = fval / fder + p = p0 - newton_step + if abs(p - p0) < tol: + return p + p0 = p + return nan + + +@hjit("f(f)") +def D_to_nu_hf(D): r"""True anomaly from parabolic anomaly. Parameters @@ -69,11 +175,20 @@ def D_to_nu(D): \nu = 2 \arctan{D} """ - return 2.0 * np.arctan(D) + return 2 * atan(D) -@jit -def nu_to_D(nu): +@vjit("f(f)") +def D_to_nu_vf(D): + """ + Vectorized D_to_nu + """ + + return D_to_nu_hf(D) + + +@hjit("f(f)") +def nu_to_D_hf(nu): r"""Parabolic anomaly from true anomaly. Parameters @@ -121,11 +236,20 @@ def nu_to_D(nu): """ # TODO: Rename to B - return np.tan(nu / 2.0) + return tan(nu / 2) -@jit -def nu_to_E(nu, ecc): +@vjit("f(f)") +def nu_to_D_vf(nu): + """ + Vectorized nu_to_D + """ + + return nu_to_D_hf(nu) + + +@hjit("f(f,f)") +def nu_to_E_hf(nu, ecc): r"""Eccentric anomaly from true anomaly. .. versionadded:: 0.4.0 @@ -156,12 +280,21 @@ def nu_to_E(nu, ecc): \in (-\pi, \pi] """ - E = 2 * np.arctan(np.sqrt((1 - ecc) / (1 + ecc)) * np.tan(nu / 2)) + E = 2 * atan(sqrt((1 - ecc) / (1 + ecc)) * tan(nu / 2)) return E -@jit -def nu_to_F(nu, ecc): +@vjit("f(f,f)") +def nu_to_E_vf(nu, ecc): + """ + Vectorized nu_to_E + """ + + return nu_to_E_hf(nu, ecc) + + +@hjit("f(f,f)") +def nu_to_F_hf(nu, ecc): r"""Hyperbolic anomaly from true anomaly. Parameters @@ -192,12 +325,21 @@ def nu_to_F(nu, ecc): F = 2 \operatorname{arctanh} \left( \sqrt{\frac{e-1}{e+1}} \tan{\frac{\nu}{2}} \right) """ - F = 2 * np.arctanh(np.sqrt((ecc - 1) / (ecc + 1)) * np.tan(nu / 2)) + F = 2 * atanh(sqrt((ecc - 1) / (ecc + 1)) * tan(nu / 2)) return F -@jit -def E_to_nu(E, ecc): +@vjit("f(f,f)") +def nu_to_F_vf(nu, ecc): + """ + Vectorized nu_to_F + """ + + return nu_to_F_hf(nu, ecc) + + +@hjit("f(f,f)") +def E_to_nu_hf(E, ecc): r"""True anomaly from eccentric anomaly. .. versionadded:: 0.4.0 @@ -228,12 +370,21 @@ def E_to_nu(E, ecc): \in (-\pi, \pi] """ - nu = 2 * np.arctan(np.sqrt((1 + ecc) / (1 - ecc)) * np.tan(E / 2)) + nu = 2 * atan(sqrt((1 + ecc) / (1 - ecc)) * tan(E / 2)) return nu -@jit -def F_to_nu(F, ecc): +@vjit("f(f,f)") +def E_to_nu_vf(E, ecc): + """ + Vectorized E_to_nu + """ + + return E_to_nu_hf(E, ecc) + + +@hjit("f(f,f)") +def F_to_nu_hf(F, ecc): r"""True anomaly from hyperbolic anomaly. Parameters @@ -257,12 +408,21 @@ def F_to_nu(F, ecc): \in (-\pi, \pi] """ - nu = 2 * np.arctan(np.sqrt((ecc + 1) / (ecc - 1)) * np.tanh(F / 2)) + nu = 2 * atan(sqrt((ecc + 1) / (ecc - 1)) * tanh(F / 2)) return nu -@jit -def M_to_E(M, ecc): +@vjit("f(f,f)") +def F_to_nu_vf(F, ecc): + """ + Vectorized F_to_nu + """ + + return F_to_nu_hf(F, ecc) + + +@hjit("f(f,f)") +def M_to_E_hf(M, ecc): """Eccentric anomaly from mean anomaly. .. versionadded:: 0.4.0 @@ -284,16 +444,25 @@ def M_to_E(M, ecc): This uses a Newton iteration on the Kepler equation. """ - if -np.pi < M < 0 or np.pi < M: + if -pi < M < 0 or pi < M: E0 = M - ecc else: E0 = M + ecc - E = _newton_elliptic(E0, args=(M, ecc)) + E = _newton_elliptic_hf(E0, M, ecc, _TOL, 50) return E -@jit -def M_to_F(M, ecc): +@vjit("f(f,f)") +def M_to_E_vf(M, ecc): + """ + Vectorized M_to_E + """ + + return M_to_E_hf(M, ecc) + + +@hjit("f(f,f)") +def M_to_F_hf(M, ecc): """Hyperbolic anomaly from mean anomaly. Parameters @@ -313,13 +482,22 @@ def M_to_F(M, ecc): This uses a Newton iteration on the hyperbolic Kepler equation. """ - F0 = np.arcsinh(M / ecc) - F = _newton_hyperbolic(F0, args=(M, ecc), maxiter=100) + F0 = asinh(M / ecc) + F = _newton_hyperbolic_hf(F0, M, ecc, _TOL, 100) return F -@jit -def M_to_D(M): +@vjit("f(f,f)") +def M_to_F_vf(M, ecc): + """ + Vectorized M_to_F + """ + + return M_to_F_hf(M, ecc) + + +@hjit("f(f)") +def M_to_D_hf(M): """Parabolic anomaly from mean anomaly. Parameters @@ -343,76 +521,17 @@ def M_to_D(M): return D -@jit -def E_to_M(E, ecc): - r"""Mean anomaly from eccentric anomaly. - - .. versionadded:: 0.4.0 - - Parameters - ---------- - E : float - Eccentric anomaly in radians. - ecc : float - Eccentricity. - - Returns - ------- - M : float - Mean anomaly. - - Warnings - -------- - The mean anomaly will be outside of (-π, π] - if the eccentric anomaly is. - No validation or wrapping is performed. - - Notes - ----- - The implementation uses the plain original Kepler equation: - - .. math:: - M = E - e \sin{E} - +@vjit("f(f)") +def M_to_D_vf(M): """ - M = E - ecc * np.sin(E) - return M - - -@jit -def F_to_M(F, ecc): - r"""Mean anomaly from hyperbolic anomaly. - - Parameters - ---------- - F : float - Hyperbolic anomaly. - ecc : float - Eccentricity (>1). - - Returns - ------- - M : float - Mean anomaly. - - Notes - ----- - As noted in [5]_, by manipulating - the parametric equations of the hyperbola - we can derive a quantity that is equivalent - to the mean anomaly in the elliptic case: - - .. math:: - - M = e \sinh{F} - F - + Vectorized M_to_D """ - M = ecc * np.sinh(F) - F - return M + return M_to_D_hf(M) -@jit -def D_to_M(D): + +@hjit("f(f)") +def D_to_M_hf(D): r"""Mean anomaly from parabolic anomaly. Parameters @@ -444,8 +563,17 @@ def D_to_M(D): return M -@jit -def fp_angle(nu, ecc): +@vjit("f(f)") +def D_to_M_vf(D): + """ + Vectorized D_to_M + """ + + return D_to_M_hf(D) + + +@hjit("f(f,f)") +def fp_angle_hf(nu, ecc): r"""Returns the flight path angle. Parameters @@ -469,4 +597,13 @@ def fp_angle(nu, ecc): \phi = \arctan(\frac {e \sin{\nu}}{1 + e \cos{\nu}}) """ - return np.arctan2(ecc * np.sin(nu), 1 + ecc * np.cos(nu)) + return atan2(ecc * sin(nu), 1 + ecc * cos(nu)) + + +@vjit("f(f,f)") +def fp_angle_vf(nu, ecc): + """ + Vectorized fp_angle + """ + + return fp_angle_hf(nu, ecc) From d7d981d036d4e58f35cde808fddcf62710acc5d8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 12:17:19 +0100 Subject: [PATCH 027/346] change to public --- src/hapsira/core/angles.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/hapsira/core/angles.py b/src/hapsira/core/angles.py index d9f0b921d..87c6178a0 100644 --- a/src/hapsira/core/angles.py +++ b/src/hapsira/core/angles.py @@ -107,30 +107,30 @@ def F_to_M_vf(F, ecc): @hjit("f(f,f,f)") -def _kepler_equation_hf(E, M, ecc): +def kepler_equation_hf(E, M, ecc): return E_to_M_hf(E, ecc) - M @hjit("f(f,f,f)") -def _kepler_equation_prime_hf(E, M, ecc): +def kepler_equation_prime_hf(E, M, ecc): return 1 - ecc * cos(E) @hjit("f(f,f,f)") -def _kepler_equation_hyper_hf(F, M, ecc): +def kepler_equation_hyper_hf(F, M, ecc): return F_to_M_hf(F, ecc) - M @hjit("f(f,f,f)") -def _kepler_equation_prime_hyper_hf(F, M, ecc): +def kepler_equation_prime_hyper_hf(F, M, ecc): return ecc * cosh(F) - 1 @hjit("f(f,f,f,f,i64)") def _newton_elliptic_hf(p0, M, ecc, tol, maxiter): for _ in range(maxiter): - fval = _kepler_equation_hf(p0, M, ecc) - fder = _kepler_equation_prime_hf(p0, M, ecc) + fval = kepler_equation_hf(p0, M, ecc) + fder = kepler_equation_prime_hf(p0, M, ecc) newton_step = fval / fder p = p0 - newton_step if abs(p - p0) < tol: @@ -142,8 +142,8 @@ def _newton_elliptic_hf(p0, M, ecc, tol, maxiter): @hjit("f(f,f,f,f,i64)") def _newton_hyperbolic_hf(p0, M, ecc, tol, maxiter): for _ in range(maxiter): - fval = _kepler_equation_hyper_hf(p0, M, ecc) - fder = _kepler_equation_prime_hyper_hf(p0, M, ecc) + fval = kepler_equation_hyper_hf(p0, M, ecc) + fder = kepler_equation_prime_hyper_hf(p0, M, ecc) newton_step = fval / fder p = p0 - newton_step if abs(p - p0) < tol: From e144b1eba0c6afc11c9179abfb29e4c7c75f02e9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 12:18:38 +0100 Subject: [PATCH 028/346] new angle helpers --- src/hapsira/core/propagation/danby.py | 6 +- src/hapsira/core/propagation/farnocchia.py | 78 +++++++++++----------- src/hapsira/core/propagation/gooding.py | 6 +- src/hapsira/core/propagation/markley.py | 18 ++--- src/hapsira/core/propagation/mikkola.py | 24 +++---- src/hapsira/core/propagation/pimienta.py | 6 +- src/hapsira/core/propagation/recseries.py | 6 +- 7 files changed, 72 insertions(+), 72 deletions(-) diff --git a/src/hapsira/core/propagation/danby.py b/src/hapsira/core/propagation/danby.py index f6a71c2cc..0f685b12b 100644 --- a/src/hapsira/core/propagation/danby.py +++ b/src/hapsira/core/propagation/danby.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from hapsira.core.angles import E_to_M, F_to_M, nu_to_E, nu_to_F +from hapsira.core.angles import E_to_M_hf, F_to_M_hf, nu_to_E_hf, nu_to_F_hf from hapsira.core.elements import coe2rv, rv2coe @@ -19,14 +19,14 @@ def danby_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter=20, rtol=1e-8): elif ecc < 1.0: # For elliptical orbit - M0 = E_to_M(nu_to_E(nu, ecc), ecc) + M0 = E_to_M_hf(nu_to_E_hf(nu, ecc), ecc) M = M0 + n * tof xma = M - 2 * np.pi * np.floor(M / 2 / np.pi) E = xma + 0.85 * np.sign(np.sin(xma)) * ecc else: # For parabolic and hyperbolic - M0 = F_to_M(nu_to_F(nu, ecc), ecc) + M0 = F_to_M_hf(nu_to_F_hf(nu, ecc), ecc) M = M0 + n * tof xma = M - 2 * np.pi * np.floor(M / 2 / np.pi) E = np.log(2 * xma / ecc + 1.8) diff --git a/src/hapsira/core/propagation/farnocchia.py b/src/hapsira/core/propagation/farnocchia.py index d20e7124b..dc19dc2bd 100644 --- a/src/hapsira/core/propagation/farnocchia.py +++ b/src/hapsira/core/propagation/farnocchia.py @@ -2,18 +2,18 @@ import numpy as np from hapsira.core.angles import ( - D_to_M, - D_to_nu, - E_to_M, - E_to_nu, - F_to_M, - F_to_nu, - M_to_D, - M_to_E, - M_to_F, - nu_to_D, - nu_to_E, - nu_to_F, + D_to_M_hf, + D_to_nu_hf, + E_to_M_hf, + E_to_nu_hf, + F_to_M_hf, + F_to_nu_hf, + M_to_D_hf, + M_to_E_hf, + M_to_F_hf, + nu_to_D_hf, + nu_to_E_hf, + nu_to_F_hf, ) from hapsira.core.elements import coe2rv, rv2coe @@ -110,7 +110,7 @@ def M_to_D_near_parabolic(M, ecc, tol=1.48e-08, maxiter=50): Parabolic eccentric anomaly. """ - D0 = M_to_D(M) + D0 = M_to_D_hf(M) for _ in range(maxiter): fval = _kepler_equation_near_parabolic(D0, M, ecc) @@ -152,18 +152,18 @@ def delta_t_from_nu(nu, ecc, k=1.0, q=1.0, delta=1e-2): assert -np.pi <= nu < np.pi if ecc < 1 - delta: # Strong elliptic - E = nu_to_E(nu, ecc) # (-pi, pi] - M = E_to_M(E, ecc) # (-pi, pi] + E = nu_to_E_hf(nu, ecc) # (-pi, pi] + M = E_to_M_hf(E, ecc) # (-pi, pi] n = np.sqrt(k * (1 - ecc) ** 3 / q**3) elif 1 - delta <= ecc < 1: - E = nu_to_E(nu, ecc) # (-pi, pi] + E = nu_to_E_hf(nu, ecc) # (-pi, pi] if delta <= 1 - ecc * np.cos(E): # Strong elliptic - M = E_to_M(E, ecc) # (-pi, pi] + M = E_to_M_hf(E, ecc) # (-pi, pi] n = np.sqrt(k * (1 - ecc) ** 3 / q**3) else: # Near parabolic - D = nu_to_D(nu) # (-∞, ∞) + D = nu_to_D_hf(nu) # (-∞, ∞) # If |nu| is far from pi this result is bounded # because the near parabolic region shrinks in its vicinity, # otherwise the eccentricity is very close to 1 @@ -172,8 +172,8 @@ def delta_t_from_nu(nu, ecc, k=1.0, q=1.0, delta=1e-2): n = np.sqrt(k / (2 * q**3)) elif ecc == 1: # Parabolic - D = nu_to_D(nu) # (-∞, ∞) - M = D_to_M(D) # (-∞, ∞) + D = nu_to_D_hf(nu) # (-∞, ∞) + M = D_to_M_hf(D) # (-∞, ∞) n = np.sqrt(k / (2 * q**3)) elif 1 + ecc * np.cos(nu) < 0: # Unfeasible region @@ -182,20 +182,20 @@ def delta_t_from_nu(nu, ecc, k=1.0, q=1.0, delta=1e-2): # NOTE: Do we need to wrap nu here? # For hyperbolic orbits, it should anyway be in # (-arccos(-1 / ecc), +arccos(-1 / ecc)) - F = nu_to_F(nu, ecc) # (-∞, ∞) + F = nu_to_F_hf(nu, ecc) # (-∞, ∞) if delta <= ecc * np.cosh(F) - 1: # Strong hyperbolic - M = F_to_M(F, ecc) # (-∞, ∞) + M = F_to_M_hf(F, ecc) # (-∞, ∞) n = np.sqrt(k * (ecc - 1) ** 3 / q**3) else: # Near parabolic - D = nu_to_D(nu) # (-∞, ∞) + D = nu_to_D_hf(nu) # (-∞, ∞) M = D_to_M_near_parabolic(D, ecc) # (-∞, ∞) n = np.sqrt(k / (2 * q**3)) elif 1 + delta < ecc: # Strong hyperbolic - F = nu_to_F(nu, ecc) # (-∞, ∞) - M = F_to_M(F, ecc) # (-∞, ∞) + F = nu_to_F_hf(nu, ecc) # (-∞, ∞) + M = F_to_M_hf(F, ecc) # (-∞, ∞) n = np.sqrt(k * (ecc - 1) ** 3 / q**3) else: raise RuntimeError @@ -232,8 +232,8 @@ def nu_from_delta_t(delta_t, ecc, k=1.0, q=1.0, delta=1e-2): M = n * delta_t # This might represent several revolutions, # so we wrap the true anomaly - E = M_to_E((M + np.pi) % (2 * np.pi) - np.pi, ecc) - nu = E_to_nu(E, ecc) + E = M_to_E_hf((M + np.pi) % (2 * np.pi) - np.pi, ecc) + nu = E_to_nu_hf(E, ecc) elif 1 - delta <= ecc < 1: E_delta = np.arccos((1 - delta) / ecc) # We compute M assuming we are in the strong elliptic case @@ -241,24 +241,24 @@ def nu_from_delta_t(delta_t, ecc, k=1.0, q=1.0, delta=1e-2): n = np.sqrt(k * (1 - ecc) ** 3 / q**3) M = n * delta_t # We check against abs(M) because E_delta could also be negative - if E_to_M(E_delta, ecc) <= abs(M): + if E_to_M_hf(E_delta, ecc) <= abs(M): # Strong elliptic, proceed # This might represent several revolutions, # so we wrap the true anomaly - E = M_to_E((M + np.pi) % (2 * np.pi) - np.pi, ecc) - nu = E_to_nu(E, ecc) + E = M_to_E_hf((M + np.pi) % (2 * np.pi) - np.pi, ecc) + nu = E_to_nu_hf(E, ecc) else: # Near parabolic, recompute M n = np.sqrt(k / (2 * q**3)) M = n * delta_t D = M_to_D_near_parabolic(M, ecc) - nu = D_to_nu(D) + nu = D_to_nu_hf(D) elif ecc == 1: # Parabolic n = np.sqrt(k / (2 * q**3)) M = n * delta_t - D = M_to_D(M) - nu = D_to_nu(D) + D = M_to_D_hf(M) + nu = D_to_nu_hf(D) elif 1 < ecc <= 1 + delta: F_delta = np.arccosh((1 + delta) / ecc) # We compute M assuming we are in the strong hyperbolic case @@ -266,23 +266,23 @@ def nu_from_delta_t(delta_t, ecc, k=1.0, q=1.0, delta=1e-2): n = np.sqrt(k * (ecc - 1) ** 3 / q**3) M = n * delta_t # We check against abs(M) because F_delta could also be negative - if F_to_M(F_delta, ecc) <= abs(M): + if F_to_M_hf(F_delta, ecc) <= abs(M): # Strong hyperbolic, proceed - F = M_to_F(M, ecc) - nu = F_to_nu(F, ecc) + F = M_to_F_hf(M, ecc) + nu = F_to_nu_hf(F, ecc) else: # Near parabolic, recompute M n = np.sqrt(k / (2 * q**3)) M = n * delta_t D = M_to_D_near_parabolic(M, ecc) - nu = D_to_nu(D) + nu = D_to_nu_hf(D) # elif 1 + delta < ecc: else: # Strong hyperbolic n = np.sqrt(k * (ecc - 1) ** 3 / q**3) M = n * delta_t - F = M_to_F(M, ecc) - nu = F_to_nu(F, ecc) + F = M_to_F_hf(M, ecc) + nu = F_to_nu_hf(F, ecc) return nu diff --git a/src/hapsira/core/propagation/gooding.py b/src/hapsira/core/propagation/gooding.py index 59f53bc0b..61c737a7f 100644 --- a/src/hapsira/core/propagation/gooding.py +++ b/src/hapsira/core/propagation/gooding.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from hapsira.core.angles import E_to_M, E_to_nu, nu_to_E +from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf from hapsira.core.elements import coe2rv, rv2coe @@ -13,7 +13,7 @@ def gooding_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter=150, rtol=1e-8): "Parabolic/Hyperbolic cases still not implemented in gooding." ) - M0 = E_to_M(nu_to_E(nu, ecc), ecc) + M0 = E_to_M_hf(nu_to_E_hf(nu, ecc), ecc) semi_axis_a = p / (1 - ecc**2) n = np.sqrt(k / np.abs(semi_axis_a) ** 3) M = M0 + n * tof @@ -34,7 +34,7 @@ def gooding_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter=150, rtol=1e-8): n += 1 E = M + psi - return E_to_nu(E, ecc) + return E_to_nu_hf(E, ecc) @jit diff --git a/src/hapsira/core/propagation/markley.py b/src/hapsira/core/propagation/markley.py index 2cedb5352..32ec6bd0c 100644 --- a/src/hapsira/core/propagation/markley.py +++ b/src/hapsira/core/propagation/markley.py @@ -2,18 +2,18 @@ import numpy as np from hapsira.core.angles import ( - E_to_M, - E_to_nu, - _kepler_equation, - _kepler_equation_prime, - nu_to_E, + E_to_M_hf, + E_to_nu_hf, + kepler_equation_hf, + kepler_equation_prime_hf, + nu_to_E_hf, ) from hapsira.core.elements import coe2rv, rv2coe @jit def markley_coe(k, p, ecc, inc, raan, argp, nu, tof): - M0 = E_to_M(nu_to_E(nu, ecc), ecc) + M0 = E_to_M_hf(nu_to_E_hf(nu, ecc), ecc) a = p / (1 - ecc**2) n = np.sqrt(k / a**3) M = M0 + n * tof @@ -40,8 +40,8 @@ def markley_coe(k, p, ecc, inc, raan, argp, nu, tof): E = (2 * r * w / (w**2 + w * q + q**2) + M) / d # Equation (26) - f0 = _kepler_equation(E, M, ecc) - f1 = _kepler_equation_prime(E, M, ecc) + f0 = kepler_equation_hf(E, M, ecc) + f1 = kepler_equation_prime_hf(E, M, ecc) f2 = ecc * np.sin(E) f3 = ecc * np.cos(E) f4 = -f2 @@ -54,7 +54,7 @@ def markley_coe(k, p, ecc, inc, raan, argp, nu, tof): ) E += delta5 - nu = E_to_nu(E, ecc) + nu = E_to_nu_hf(E, ecc) return nu diff --git a/src/hapsira/core/propagation/mikkola.py b/src/hapsira/core/propagation/mikkola.py index f40d18412..5daa5acb8 100644 --- a/src/hapsira/core/propagation/mikkola.py +++ b/src/hapsira/core/propagation/mikkola.py @@ -2,13 +2,13 @@ import numpy as np from hapsira.core.angles import ( - D_to_nu, - E_to_M, - E_to_nu, - F_to_M, - F_to_nu, - nu_to_E, - nu_to_F, + D_to_nu_hf, + E_to_M_hf, + E_to_nu_hf, + F_to_M_hf, + F_to_nu_hf, + nu_to_E_hf, + nu_to_F_hf, ) from hapsira.core.elements import coe2rv, rv2coe @@ -22,10 +22,10 @@ def mikkola_coe(k, p, ecc, inc, raan, argp, nu, tof): if ecc < 1.0: # Equation (9a) alpha = (1 - ecc) / (4 * ecc + 1 / 2) - M0 = E_to_M(nu_to_E(nu, ecc), ecc) + M0 = E_to_M_hf(nu_to_E_hf(nu, ecc), ecc) else: alpha = (ecc - 1) / (4 * ecc + 1 / 2) - M0 = F_to_M(nu_to_F(nu, ecc), ecc) + M0 = F_to_M_hf(nu_to_F_hf(nu, ecc), ecc) M = M0 + n * tof beta = M / 2 / (4 * ecc + 1 / 2) @@ -82,14 +82,14 @@ def mikkola_coe(k, p, ecc, inc, raan, argp, nu, tof): E += u5 if ecc < 1.0: - nu = E_to_nu(E, ecc) + nu = E_to_nu_hf(E, ecc) else: if ecc == 1.0: # Parabolic - nu = D_to_nu(E) + nu = D_to_nu_hf(E) else: # Hyperbolic - nu = F_to_nu(E, ecc) + nu = F_to_nu_hf(E, ecc) return nu diff --git a/src/hapsira/core/propagation/pimienta.py b/src/hapsira/core/propagation/pimienta.py index 3914466ad..13e2c492b 100644 --- a/src/hapsira/core/propagation/pimienta.py +++ b/src/hapsira/core/propagation/pimienta.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from hapsira.core.angles import E_to_M, E_to_nu, nu_to_E +from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf from hapsira.core.elements import coe2rv, rv2coe @@ -11,7 +11,7 @@ def pimienta_coe(k, p, ecc, inc, raan, argp, nu, tof): # TODO: Do something to allow parabolic and hyperbolic orbits? n = np.sqrt(k * (1 - ecc) ** 3 / q**3) - M0 = E_to_M(nu_to_E(nu, ecc), ecc) + M0 = E_to_M_hf(nu_to_E_hf(nu, ecc), ecc) M = M0 + n * tof @@ -333,7 +333,7 @@ def pimienta_coe(k, p, ecc, inc, raan, argp, nu, tof): + 15 * w ) - return E_to_nu(E, ecc) + return E_to_nu_hf(E, ecc) @jit diff --git a/src/hapsira/core/propagation/recseries.py b/src/hapsira/core/propagation/recseries.py index ae99feba2..55d818dd6 100644 --- a/src/hapsira/core/propagation/recseries.py +++ b/src/hapsira/core/propagation/recseries.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from hapsira.core.angles import E_to_M, E_to_nu, nu_to_E +from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf from hapsira.core.elements import coe2rv, rv2coe @@ -41,7 +41,7 @@ def recseries_coe( # Solving for elliptical orbit # compute initial mean anoamly - M0 = E_to_M(nu_to_E(nu, ecc), ecc) + M0 = E_to_M_hf(nu_to_E_hf(nu, ecc), ecc) # final mean anaomaly M = M0 + n * tof # snapping anomaly to [0,pi] range @@ -64,7 +64,7 @@ def recseries_coe( break E = En - return E_to_nu(E, ecc) + return E_to_nu_hf(E, ecc) else: # Parabolic/Hyperbolic orbits are not supported From 4eabbc89b9c91c25a6bcc026392c0ea837cb09b4 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 12:34:09 +0100 Subject: [PATCH 029/346] fix --- src/hapsira/core/angles.py | 4 ++-- src/hapsira/core/elements.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/hapsira/core/angles.py b/src/hapsira/core/angles.py index 87c6178a0..5051b054f 100644 --- a/src/hapsira/core/angles.py +++ b/src/hapsira/core/angles.py @@ -126,7 +126,7 @@ def kepler_equation_prime_hyper_hf(F, M, ecc): return ecc * cosh(F) - 1 -@hjit("f(f,f,f,f,i64)") +@hjit("f(f,f,f,f,i8)") def _newton_elliptic_hf(p0, M, ecc, tol, maxiter): for _ in range(maxiter): fval = kepler_equation_hf(p0, M, ecc) @@ -139,7 +139,7 @@ def _newton_elliptic_hf(p0, M, ecc, tol, maxiter): return nan -@hjit("f(f,f,f,f,i64)") +@hjit("f(f,f,f,f,i8)") def _newton_hyperbolic_hf(p0, M, ecc, tol, maxiter): for _ in range(maxiter): fval = kepler_equation_hyper_hf(p0, M, ecc) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 67f12d5e3..2c055ba1f 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -8,7 +8,7 @@ import numpy as np from numpy import cos, cross, sin, sqrt -from hapsira.core.angles import E_to_nu, F_to_nu +from hapsira.core.angles import E_to_nu_hf, F_to_nu_hf from hapsira.core.util import rotation_matrix from .math.linalg import norm @@ -408,11 +408,11 @@ def rv2coe(k, r, v, tol=1e-8): if a > 0: e_se = (r @ v) / sqrt(ka) e_ce = norm(r) * (v @ v) / k - 1 - nu = E_to_nu(np.arctan2(e_se, e_ce), ecc) + nu = E_to_nu_hf(np.arctan2(e_se, e_ce), ecc) else: e_sh = (r @ v) / sqrt(-ka) e_ch = norm(r) * (norm(v) ** 2) / k - 1 - nu = F_to_nu(np.log((e_ch + e_sh) / (e_ch - e_sh)) / 2, ecc) + nu = F_to_nu_hf(np.log((e_ch + e_sh) / (e_ch - e_sh)) / 2, ecc) raan = np.arctan2(n[1], n[0]) % (2 * np.pi) px = r @ n From ba6e3c62676c505153bd26ddb8ebb8e9fb467ae2 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 12:34:20 +0100 Subject: [PATCH 030/346] use new api --- src/hapsira/twobody/angles.py | 52 +++++++++++++++++------------------ 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/src/hapsira/twobody/angles.py b/src/hapsira/twobody/angles.py index a07091be4..1cf38326f 100644 --- a/src/hapsira/twobody/angles.py +++ b/src/hapsira/twobody/angles.py @@ -2,19 +2,19 @@ from astropy import units as u from hapsira.core.angles import ( - D_to_M as D_to_M_fast, - D_to_nu as D_to_nu_fast, - E_to_M as E_to_M_fast, - E_to_nu as E_to_nu_fast, - F_to_M as F_to_M_fast, - F_to_nu as F_to_nu_fast, - M_to_D as M_to_D_fast, - M_to_E as M_to_E_fast, - M_to_F as M_to_F_fast, - fp_angle as fp_angle_fast, - nu_to_D as nu_to_D_fast, - nu_to_E as nu_to_E_fast, - nu_to_F as nu_to_F_fast, + D_to_M_vf, + D_to_nu_vf, + E_to_M_vf, + E_to_nu_vf, + F_to_M_vf, + F_to_nu_vf, + M_to_D_vf, + M_to_E_vf, + M_to_F_vf, + fp_angle_vf, + nu_to_D_vf, + nu_to_E_vf, + nu_to_F_vf, ) @@ -38,7 +38,7 @@ def D_to_nu(D): "Robust resolution of Kepler’s equation in all eccentricity regimes." Celestial Mechanics and Dynamical Astronomy 116, no. 1 (2013): 21-34. """ - return (D_to_nu_fast(D.to_value(u.rad)) * u.rad).to(D.unit) + return (D_to_nu_vf(D.to_value(u.rad)) * u.rad).to(D.unit) @u.quantity_input(nu=u.rad) @@ -61,7 +61,7 @@ def nu_to_D(nu): "Robust resolution of Kepler’s equation in all eccentricity regimes." Celestial Mechanics and Dynamical Astronomy 116, no. 1 (2013): 21-34. """ - return (nu_to_D_fast(nu.to_value(u.rad)) * u.rad).to(nu.unit) + return (nu_to_D_vf(nu.to_value(u.rad)) * u.rad).to(nu.unit) @u.quantity_input(nu=u.rad, ecc=u.one) @@ -83,7 +83,7 @@ def nu_to_E(nu, ecc): Eccentric anomaly. """ - return (nu_to_E_fast(nu.to_value(u.rad), ecc.value) * u.rad).to(nu.unit) + return (nu_to_E_vf(nu.to_value(u.rad), ecc.value) * u.rad).to(nu.unit) @u.quantity_input(nu=u.rad, ecc=u.one) @@ -107,7 +107,7 @@ def nu_to_F(nu, ecc): Taken from Curtis, H. (2013). *Orbital mechanics for engineering students*. 167 """ - return (nu_to_F_fast(nu.to_value(u.rad), ecc.value) * u.rad).to(nu.unit) + return (nu_to_F_vf(nu.to_value(u.rad), ecc.value) * u.rad).to(nu.unit) @u.quantity_input(E=u.rad, ecc=u.one) @@ -129,7 +129,7 @@ def E_to_nu(E, ecc): True anomaly. """ - return (E_to_nu_fast(E.to_value(u.rad), ecc.value) * u.rad).to(E.unit) + return (E_to_nu_vf(E.to_value(u.rad), ecc.value) * u.rad).to(E.unit) @u.quantity_input(F=u.rad, ecc=u.one) @@ -149,7 +149,7 @@ def F_to_nu(F, ecc): True anomaly. """ - return (F_to_nu_fast(F.to_value(u.rad), ecc.value) * u.rad).to(F.unit) + return (F_to_nu_vf(F.to_value(u.rad), ecc.value) * u.rad).to(F.unit) @u.quantity_input(M=u.rad, ecc=u.one) @@ -171,7 +171,7 @@ def M_to_E(M, ecc): Eccentric anomaly. """ - return (M_to_E_fast(M.to_value(u.rad), ecc.value) * u.rad).to(M.unit) + return (M_to_E_vf(M.to_value(u.rad), ecc.value) * u.rad).to(M.unit) @u.quantity_input(M=u.rad, ecc=u.one) @@ -191,7 +191,7 @@ def M_to_F(M, ecc): Hyperbolic eccentric anomaly. """ - return (M_to_F_fast(M.to_value(u.rad), ecc.value) * u.rad).to(M.unit) + return (M_to_F_vf(M.to_value(u.rad), ecc.value) * u.rad).to(M.unit) @u.quantity_input(M=u.rad, ecc=u.one) @@ -209,7 +209,7 @@ def M_to_D(M): Parabolic eccentric anomaly. """ - return (M_to_D_fast(M.to_value(u.rad)) * u.rad).to(M.unit) + return (M_to_D_vf(M.to_value(u.rad)) * u.rad).to(M.unit) @u.quantity_input(E=u.rad, ecc=u.one) @@ -231,7 +231,7 @@ def E_to_M(E, ecc): Mean anomaly. """ - return (E_to_M_fast(E.to_value(u.rad), ecc.value) * u.rad).to(E.unit) + return (E_to_M_vf(E.to_value(u.rad), ecc.value) * u.rad).to(E.unit) @u.quantity_input(F=u.rad, ecc=u.one) @@ -251,7 +251,7 @@ def F_to_M(F, ecc): Mean anomaly. """ - return (F_to_M_fast(F.to_value(u.rad), ecc.value) * u.rad).to(F.unit) + return (F_to_M_vf(F.to_value(u.rad), ecc.value) * u.rad).to(F.unit) @u.quantity_input(D=u.rad, ecc=u.one) @@ -269,7 +269,7 @@ def D_to_M(D): Mean anomaly. """ - return (D_to_M_fast(D.to_value(u.rad)) * u.rad).to(D.unit) + return (D_to_M_vf(D.to_value(u.rad)) * u.rad).to(D.unit) @u.quantity_input(nu=u.rad, ecc=u.one) @@ -290,4 +290,4 @@ def fp_angle(nu, ecc): Algorithm taken from Vallado 2007, pp. 113. """ - return (fp_angle_fast(nu.to_value(u.rad), ecc.value) * u.rad).to(nu.unit) + return (fp_angle_vf(nu.to_value(u.rad), ecc.value) * u.rad).to(nu.unit) From 4661c5e5449d2410b4fe9ff65e4f6510bc6b6951 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 18:52:03 +0100 Subject: [PATCH 031/346] python 3.8 compat --- src/hapsira/settings.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/hapsira/settings.py b/src/hapsira/settings.py index 18a1b85dd..c45cdbf79 100644 --- a/src/hapsira/settings.py +++ b/src/hapsira/settings.py @@ -1,5 +1,5 @@ import os -from typing import Any, Generator, Optional, Type +from typing import Any, Generator, Optional, Tuple, Type __all__ = [ "Setting", @@ -26,7 +26,7 @@ class Setting: Holds one setting settable by user before sub-module import """ - def __init__(self, name: str, default: Any, options: Optional[tuple[Any]] = None): + def __init__(self, name: str, default: Any, options: Optional[Tuple[Any]] = None): self._name = name self._type = type(default) self._value = default @@ -59,7 +59,7 @@ def type_(self) -> Type: return self._type @property - def options(self) -> Optional[tuple[Any]]: + def options(self) -> Optional[Tuple[Any]]: """ Return options for value """ From 5b8d21a3fd14b99f013ce2d19687498e3c45ff63 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 18:58:27 +0100 Subject: [PATCH 032/346] exports --- src/hapsira/core/angles.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/src/hapsira/core/angles.py b/src/hapsira/core/angles.py index 5051b054f..33bc04735 100644 --- a/src/hapsira/core/angles.py +++ b/src/hapsira/core/angles.py @@ -20,6 +20,40 @@ _TOL = 1.48e-08 +__all__ = [ + "E_to_M_hf", + "E_to_M_vf", + "F_to_M_hf", + "F_to_M_vf", + "kepler_equation_hf", + "kepler_equation_prime_hf", + "kepler_equation_hyper_hf", + "kepler_equation_prime_hyper_hf", + "D_to_nu_hf", + "D_to_nu_vf", + "nu_to_D_hf", + "nu_to_D_vf", + "nu_to_E_hf", + "nu_to_E_vf", + "nu_to_F_hf", + "nu_to_F_vf", + "E_to_nu_hf", + "E_to_nu_vf", + "F_to_nu_hf", + "F_to_nu_vf", + "M_to_E_hf", + "M_to_E_vf", + "M_to_F_hf", + "M_to_F_vf", + "M_to_D_hf", + "M_to_D_vf", + "D_to_M_hf", + "D_to_M_vf", + "fp_angle_hf", + "fp_angle_vf", +] + + @hjit("f(f,f)") def E_to_M_hf(E, ecc): r"""Mean anomaly from eccentric anomaly. From d3f67b988fa30dfa9c9a6a1a6ead8c4cae794415 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 19:00:21 +0100 Subject: [PATCH 033/346] spelling fix --- src/hapsira/core/elements.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 2c055ba1f..9954aa001 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -38,7 +38,7 @@ def eccentricity_vector(k, r, v): @jit def circular_velocity(k, a): - r"""Compute circular velocity for a given body given thegravitational parameter and the semimajor axis. + r"""Compute circular velocity for a given body given the gravitational parameter and the semimajor axis. .. math:: From 91913aaae5d01b35e7f5bc262fe724ae849deca1 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 19:30:08 +0100 Subject: [PATCH 034/346] allow matrix type --- src/hapsira/core/jit.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index e130547c0..97e7d5220 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -45,9 +45,10 @@ def _parse_signatures(signature: str) -> str | list[str]: ) return signature - signature = signature.replace( - "V", "Tuple([f,f,f])" - ) # TODO hope for support of "f[:]" return values in cuda target; 2D/4D vectors? + # TODO hope for support of "f[:]" return values in cuda target; 2D/4D vectors? + signature = signature.replace("M", "Tuple([V,V,V])") # matrix is a tuple of vectors + signature = signature.replace("V", "Tuple([f,f,f])") # vector is a tuple of floats + return [signature.replace("f", dtype) for dtype in PRECISIONS] From 8bc7008cf84043d232353d68983f738b7227b2f9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 19:32:15 +0100 Subject: [PATCH 035/346] prepare for guv --- src/hapsira/core/jit.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 97e7d5220..702eac15b 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -31,8 +31,15 @@ def _parse_signatures(signature: str) -> str | list[str]: Automatically generate signatures for single and double """ + if "->" in signature: # this is likely a layout for guvectorize + logger.warning( + "jit signature: likely a layout for guvectorize, not parsing (%s)", + signature, + ) + return signature + if not any( - notation in signature for notation in ("f", "V") + notation in signature for notation in ("f", "V", "M") ): # leave this signature as it is logger.warning( "jit signature: no special notation, not parsing (%s)", signature From 28cf7ac3455c0a216e9c3fb295a733fd1b7781ba Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 19:35:20 +0100 Subject: [PATCH 036/346] add guvectorize --- src/hapsira/core/jit.py | 46 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 702eac15b..6a90699c7 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -12,6 +12,7 @@ "PRECISIONS", "hjit", "vjit", + "gjit", "sjit", ] @@ -156,6 +157,51 @@ def wrapper(inner_func: Callable) -> Callable: return wrapper +def gjit(*args, **kwargs) -> Callable: + """ + General vectorize on array, pre-configured, user-facing, switches compiler targets. + Functions decorated by it can always be called directly if needed. + """ + + if len(args) == 1 and callable(args[0]): + outer_func = args[0] + args = tuple() + else: + outer_func = None + + if len(args) > 0 and isinstance(args[0], str): + args = _parse_signatures(args[0]), *args[1:] + + def wrapper(inner_func: Callable) -> Callable: + """ + Applies JIT + """ + + cfg = dict( + target=settings["TARGET"].value, + ) + if settings["TARGET"].value != "cuda": + cfg["nopython"] = settings["NOPYTHON"].value + cfg.update(kwargs) + + logger.debug( + "gjit: func=%s, args=%s, kwargs=%s", + getattr(inner_func, "__name__", repr(inner_func)), + repr(args), + repr(cfg), + ) + + return nb.guvectorize( + *args, + **cfg, + )(inner_func) + + if outer_func is not None: + return wrapper(outer_func) + + return wrapper + + def sjit(*args, **kwargs) -> Callable: """ Regular "scalar" (n)jit, pre-configured, potentially user-facing, always CPU compiler target. From 1c21034f353ac64ffa16899e678f1dfa24a17e58 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 19:42:00 +0100 Subject: [PATCH 037/346] make sure gufuncs return void --- src/hapsira/core/jit.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 6a90699c7..68ad388f1 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -27,7 +27,7 @@ PRECISIONS = ("f4", "f8") # TODO allow f2, i.e. half, for CUDA at least? -def _parse_signatures(signature: str) -> str | list[str]: +def _parse_signatures(signature: str, noreturn: bool = False) -> str | list[str]: """ Automatically generate signatures for single and double """ @@ -39,6 +39,11 @@ def _parse_signatures(signature: str) -> str | list[str]: ) return signature + if noreturn and not signature.startswith("void("): + raise JitError( + "function does not allow return values, likely compiled via guvectorize" + ) + if not any( notation in signature for notation in ("f", "V", "M") ): # leave this signature as it is @@ -170,7 +175,7 @@ def gjit(*args, **kwargs) -> Callable: outer_func = None if len(args) > 0 and isinstance(args[0], str): - args = _parse_signatures(args[0]), *args[1:] + args = _parse_signatures(args[0], noreturn=True), *args[1:] def wrapper(inner_func: Callable) -> Callable: """ From 8c247b08611a0445094bfd25b5a85a828210aa94 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 30 Dec 2023 19:43:24 +0100 Subject: [PATCH 038/346] python 3.8 compat --- src/hapsira/core/jit.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 68ad388f1..907eedef7 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -1,4 +1,4 @@ -from typing import Callable +from typing import Callable, List, Union import numba as nb from numba import cuda @@ -27,7 +27,7 @@ PRECISIONS = ("f4", "f8") # TODO allow f2, i.e. half, for CUDA at least? -def _parse_signatures(signature: str, noreturn: bool = False) -> str | list[str]: +def _parse_signatures(signature: str, noreturn: bool = False) -> Union[str, List[str]]: """ Automatically generate signatures for single and double """ From 8e5f5523d198731073a4be86640953f4d8a4d4dd Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 31 Dec 2023 23:53:49 +0100 Subject: [PATCH 039/346] arr to tup helper for refactor --- src/hapsira/core/jit.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 907eedef7..b4a2aea3c 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -249,3 +249,8 @@ def wrapper(inner_func: Callable) -> Callable: return wrapper(outer_func) return wrapper + + +@hjit("V(f[:])") # TODO remove, only for code refactor +def _arr2tup_hf(x): + return x[0], x[1], x[2] From 7166cdf98d3b02055fd23230a6b9c483e142f8a2 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 31 Dec 2023 23:55:08 +0100 Subject: [PATCH 040/346] new norm funcs; matmul vv --- src/hapsira/core/math/linalg.py | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 26b604270..28404cd3d 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -1,11 +1,25 @@ -from numba import njit as jit -import numpy as np +from math import sqrt + +from ..jit import hjit, vjit __all__ = [ - "norm", + "matmul_VV_hf", + "norm_hf", + "norm_vf", ] -@jit -def norm(arr): - return np.sqrt(arr @ arr) +@hjit("f(V,V)") +def matmul_VV_hf(a, b): + return a[0] * b[0] + a[1] * b[1] + a[2] * b[2] + + +@hjit("f(V)") +def norm_hf(a): + return sqrt(matmul_VV_hf(a, a)) + + +@vjit("f(f,f,f)") +def norm_vf(a, b, c): + # TODO add axis setting in some way for util.norm? + return norm_hf((a, b, c)) From ef8f162f6ec4c6fa30770d7d595a56de3893c662 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 31 Dec 2023 23:55:25 +0100 Subject: [PATCH 041/346] use new norm func --- src/hapsira/core/czml_utils.py | 6 +++--- src/hapsira/core/elements.py | 21 +++++++++++---------- src/hapsira/core/events.py | 9 +++++---- src/hapsira/core/flybys.py | 9 +++++---- src/hapsira/core/iod.py | 15 ++++++++++----- src/hapsira/core/maneuver.py | 17 +++++++++-------- src/hapsira/core/perturbations.py | 22 +++++++++++++--------- src/hapsira/core/propagation/vallado.py | 5 +++-- src/hapsira/core/spheroid_location.py | 11 ++++++----- src/hapsira/core/thrust/change_a_inc.py | 7 ++++--- src/hapsira/core/thrust/change_argp.py | 7 ++++--- src/hapsira/core/thrust/change_ecc_inc.py | 11 +++++++---- src/hapsira/twobody/events.py | 8 ++++---- src/hapsira/util.py | 4 ++-- 14 files changed, 86 insertions(+), 66 deletions(-) diff --git a/src/hapsira/core/czml_utils.py b/src/hapsira/core/czml_utils.py index d479c3bed..eb3b275da 100644 --- a/src/hapsira/core/czml_utils.py +++ b/src/hapsira/core/czml_utils.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from .math.linalg import norm +from .math.linalg import norm_hf @jit @@ -95,7 +95,7 @@ def project_point_on_ellipsoid(x, y, z, a, b, c): """ p1, p2 = intersection_ellipsoid_line(x, y, z, x, y, z, a, b, c) - norm_1 = norm(np.array([p1[0] - x, p1[1] - y, p1[2] - z])) - norm_2 = norm(np.array([p2[0] - x, p2[1] - y, p2[2] - z])) + norm_1 = norm_hf((p1[0] - x, p1[1] - y, p1[2] - z)) + norm_2 = norm_hf((p2[0] - x, p2[1] - y, p2[2] - z)) return p1 if norm_1 <= norm_2 else p2 diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 9954aa001..49d228fd2 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -11,7 +11,8 @@ from hapsira.core.angles import E_to_nu_hf, F_to_nu_hf from hapsira.core.util import rotation_matrix -from .math.linalg import norm +from .jit import _arr2tup_hf +from .math.linalg import norm_hf @jit @@ -33,7 +34,7 @@ def eccentricity_vector(k, r, v): v : numpy.ndarray Velocity vector (km / s) """ - return ((v @ v - k / norm(r)) * r - (r @ v) * v) / k + return ((v @ v - k / norm_hf(_arr2tup_hf(r))) * r - (r @ v) * v) / k @jit @@ -381,10 +382,10 @@ def rv2coe(k, r, v, tol=1e-8): """ h = cross(r, v) n = cross([0, 0, 1], h) - e = ((v @ v - k / norm(r)) * r - (r @ v) * v) / k - ecc = norm(e) + e = ((v @ v - k / norm_hf(_arr2tup_hf(r))) * r - (r @ v) * v) / k + ecc = norm_hf(_arr2tup_hf(e)) p = (h @ h) / k - inc = np.arccos(h[2] / norm(h)) + inc = np.arccos(h[2] / norm_hf(_arr2tup_hf(h))) circular = ecc < tol equatorial = abs(inc) < tol @@ -392,12 +393,12 @@ def rv2coe(k, r, v, tol=1e-8): if equatorial and not circular: raan = 0 argp = np.arctan2(e[1], e[0]) % (2 * np.pi) # Longitude of periapsis - nu = np.arctan2((h @ cross(e, r)) / norm(h), r @ e) + nu = np.arctan2((h @ cross(e, r)) / norm_hf(_arr2tup_hf(h)), r @ e) elif not equatorial and circular: raan = np.arctan2(n[1], n[0]) % (2 * np.pi) argp = 0 # Argument of latitude - nu = np.arctan2((r @ cross(h, n)) / norm(h), r @ n) + nu = np.arctan2((r @ cross(h, n)) / norm_hf(_arr2tup_hf(h)), r @ n) elif equatorial and circular: raan = 0 argp = 0 @@ -407,16 +408,16 @@ def rv2coe(k, r, v, tol=1e-8): ka = k * a if a > 0: e_se = (r @ v) / sqrt(ka) - e_ce = norm(r) * (v @ v) / k - 1 + e_ce = norm_hf(_arr2tup_hf(r)) * (v @ v) / k - 1 nu = E_to_nu_hf(np.arctan2(e_se, e_ce), ecc) else: e_sh = (r @ v) / sqrt(-ka) - e_ch = norm(r) * (norm(v) ** 2) / k - 1 + e_ch = norm_hf(_arr2tup_hf(r)) * (norm_hf(_arr2tup_hf(v)) ** 2) / k - 1 nu = F_to_nu_hf(np.log((e_ch + e_sh) / (e_ch - e_sh)) / 2, ecc) raan = np.arctan2(n[1], n[0]) % (2 * np.pi) px = r @ n - py = (r @ cross(h, n)) / norm(h) + py = (r @ cross(h, n)) / norm_hf(_arr2tup_hf(h)) argp = (np.arctan2(py, px) - nu) % (2 * np.pi) nu = (nu + np.pi) % (2 * np.pi) - np.pi diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index 44b736b63..c5b5a9290 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -4,7 +4,8 @@ from hapsira.core.elements import coe_rotation_matrix, rv2coe from hapsira.core.util import planetocentric_to_AltAz -from .math.linalg import norm +from .jit import _arr2tup_hf +from .math.linalg import norm_hf @jit @@ -41,7 +42,7 @@ def eclipse_function(k, u_, r_sec, R_sec, R_primary, umbra=True): # Make arrays contiguous for faster dot product with numba. P_, Q_ = np.ascontiguousarray(PQW[:, 0]), np.ascontiguousarray(PQW[:, 1]) - r_sec_norm = norm(r_sec) + r_sec_norm = norm_hf(_arr2tup_hf(r_sec)) beta = (P_ @ r_sec) / r_sec_norm zeta = (Q_ @ r_sec) / r_sec_norm @@ -78,8 +79,8 @@ def line_of_sight(r1, r2, R): located by r1 and r2, else negative. """ - r1_norm = norm(r1) - r2_norm = norm(r2) + r1_norm = norm_hf(_arr2tup_hf(r1)) + r2_norm = norm_hf(_arr2tup_hf(r2)) theta = np.arccos((r1 @ r2) / r1_norm / r2_norm) theta_1 = np.arccos(R / r1_norm) diff --git a/src/hapsira/core/flybys.py b/src/hapsira/core/flybys.py index 4849a3d44..9a90fffff 100644 --- a/src/hapsira/core/flybys.py +++ b/src/hapsira/core/flybys.py @@ -4,7 +4,8 @@ import numpy as np from numpy import cross -from .math.linalg import norm +from .jit import _arr2tup_hf +from .math.linalg import norm_hf @jit @@ -33,7 +34,7 @@ def compute_flyby(v_spacecraft, v_body, k, r_p, theta): """ v_inf_1 = v_spacecraft - v_body # Hyperbolic excess velocity - v_inf = norm(v_inf_1) + v_inf = norm_hf(_arr2tup_hf(v_inf_1)) ecc = 1 + r_p * v_inf**2 / k # Eccentricity of the entry hyperbola delta = 2 * np.arcsin(1 / ecc) # Turn angle @@ -48,7 +49,7 @@ def compute_flyby(v_spacecraft, v_body, k, r_p, theta): S_vec = v_inf_1 / v_inf c_vec = np.array([0, 0, 1]) T_vec = cross(S_vec, c_vec) - T_vec = T_vec / norm(T_vec) + T_vec = T_vec / norm_hf(_arr2tup_hf(T_vec)) R_vec = cross(S_vec, T_vec) # This vector defines the B-Plane @@ -62,7 +63,7 @@ def compute_flyby(v_spacecraft, v_body, k, r_p, theta): # And now we rotate the outbound hyperbolic excess velocity # u_vec = v_inf_1 / norm(v_inf) = S_vec v_vec = cross(rot_v, v_inf_1) - v_vec = v_vec / norm(v_vec) + v_vec = v_vec / norm_hf(_arr2tup_hf(v_vec)) v_inf_2 = v_inf * (np.cos(delta) * S_vec + np.sin(delta) * v_vec) diff --git a/src/hapsira/core/iod.py b/src/hapsira/core/iod.py index a6ed6f7e3..c3085d347 100644 --- a/src/hapsira/core/iod.py +++ b/src/hapsira/core/iod.py @@ -2,7 +2,8 @@ import numpy as np from numpy import cross, pi -from .math.linalg import norm +from .jit import _arr2tup_hf +from .math.linalg import norm_hf from .math.special import hyp2f1b, stumpff_c2 as c2, stumpff_c3 as c3 @@ -110,8 +111,8 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): t_m = 1 if prograde else -1 - norm_r0 = norm(r0) - norm_r = norm(r) + norm_r0 = norm_hf(_arr2tup_hf(r0)) + norm_r = norm_hf(_arr2tup_hf(r)) norm_r0_times_norm_r = norm_r0 * norm_r norm_r0_plus_norm_r = norm_r0 + norm_r @@ -213,7 +214,11 @@ def izzo(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol): # Chord c = r2 - r1 - c_norm, r1_norm, r2_norm = norm(c), norm(r1), norm(r2) + c_norm, r1_norm, r2_norm = ( + norm_hf(_arr2tup_hf(c)), + norm_hf(_arr2tup_hf(r1)), + norm_hf(_arr2tup_hf(r2)), + ) # Semiperimeter s = (r1_norm + r2_norm + c_norm) * 0.5 @@ -221,7 +226,7 @@ def izzo(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol): # Versors i_r1, i_r2 = r1 / r1_norm, r2 / r2_norm i_h = cross(i_r1, i_r2) - i_h = i_h / norm(i_h) # Fixed from paper + i_h = i_h / norm_hf(_arr2tup_hf(i_h)) # Fixed from paper # Geometry of the problem ll = np.sqrt(1 - min(1.0, c_norm / s)) diff --git a/src/hapsira/core/maneuver.py b/src/hapsira/core/maneuver.py index b7f7936f0..2aafc4e11 100644 --- a/src/hapsira/core/maneuver.py +++ b/src/hapsira/core/maneuver.py @@ -6,7 +6,8 @@ from hapsira.core.elements import coe_rotation_matrix, rv2coe, rv_pqw -from .math.linalg import norm +from .jit import _arr2tup_hf +from .math.linalg import norm_hf @jit @@ -42,13 +43,13 @@ def hohmann(k, rv, r_f): """ _, ecc, inc, raan, argp, nu = rv2coe(k, *rv) - h_i = norm(cross(*rv)) + h_i = norm_hf(_arr2tup_hf(cross(*rv))) p_i = h_i**2 / k r_i, v_i = rv_pqw(k, p_i, ecc, nu) - r_i = norm(r_i) - v_i = norm(v_i) + r_i = norm_hf(_arr2tup_hf(r_i)) + v_i = norm_hf(_arr2tup_hf(v_i)) a_trans = (r_i + r_f) / 2 dv_a = np.sqrt(2 * k / r_i - k / a_trans) - v_i @@ -112,13 +113,13 @@ def bielliptic(k, r_b, r_f, rv): """ _, ecc, inc, raan, argp, nu = rv2coe(k, *rv) - h_i = norm(cross(*rv)) + h_i = norm_hf(_arr2tup_hf(cross(*rv))) p_i = h_i**2 / k r_i, v_i = rv_pqw(k, p_i, ecc, nu) - r_i = norm(r_i) - v_i = norm(v_i) + r_i = norm_hf(_arr2tup_hf(r_i)) + v_i = norm_hf(_arr2tup_hf(v_i)) a_trans1 = (r_i + r_b) / 2 a_trans2 = (r_b + r_f) / 2 @@ -191,6 +192,6 @@ def correct_pericenter(k, R, J2, max_delta_r, v, a, inc, ecc): delta_t = abs(delta_w / dw) delta_v = 0.5 * n * a * ecc * abs(delta_w) - vf_ = v / norm(v) * delta_v + vf_ = v / norm_hf(_arr2tup_hf(v)) * delta_v return delta_t, vf_ diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index 3b1800d73..df621ef42 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -3,7 +3,8 @@ from hapsira.core.events import line_of_sight as line_of_sight_fast -from .math.linalg import norm +from .jit import _arr2tup_hf +from .math.linalg import norm_hf @jit @@ -36,7 +37,7 @@ def J2_perturbation(t0, state, k, J2, R): """ r_vec = state[:3] - r = norm(r_vec) + r = norm_hf(_arr2tup_hf(r_vec)) factor = (3.0 / 2.0) * k * J2 * (R**2) / (r**5) @@ -71,7 +72,7 @@ def J3_perturbation(t0, state, k, J3, R): """ r_vec = state[:3] - r = norm(r_vec) + r = norm_hf(_arr2tup_hf(r_vec)) factor = (1.0 / 2.0) * k * J3 * (R**3) / (r**5) cos_phi = r_vec[2] / r @@ -119,10 +120,10 @@ def atmospheric_drag_exponential(t0, state, k, R, C_D, A_over_m, H0, rho0): the atmospheric density model is rho(H) = rho0 x exp(-H / H0) """ - H = norm(state[:3]) + H = norm_hf(_arr2tup_hf(state[:3])) v_vec = state[3:] - v = norm(v_vec) + v = norm_hf(_arr2tup_hf(v_vec)) B = C_D * A_over_m rho = rho0 * np.exp(-(H - R) / H0) @@ -161,7 +162,7 @@ def atmospheric_drag(t0, state, k, C_D, A_over_m, rho): """ v_vec = state[3:] - v = norm(v_vec) + v = norm_hf(_arr2tup_hf(v_vec)) B = C_D * A_over_m return -(1.0 / 2.0) * rho * B * v * v_vec @@ -196,7 +197,10 @@ def third_body(t0, state, k, k_third, perturbation_body): """ body_r = perturbation_body(t0) delta_r = body_r - state[:3] - return k_third * delta_r / norm(delta_r) ** 3 - k_third * body_r / norm(body_r) ** 3 + return ( + k_third * delta_r / norm_hf(_arr2tup_hf(delta_r)) ** 3 + - k_third * body_r / norm_hf(_arr2tup_hf(body_r)) ** 3 + ) def radiation_pressure(t0, state, k, R, C_R, A_over_m, Wdivc_s, star): @@ -234,7 +238,7 @@ def radiation_pressure(t0, state, k, R, C_R, A_over_m, Wdivc_s, star): """ r_star = star(t0) r_sat = state[:3] - P_s = Wdivc_s / (norm(r_star) ** 2) + P_s = Wdivc_s / (norm_hf(_arr2tup_hf(r_star)) ** 2) nu = float(line_of_sight_fast(r_sat, r_star, R) > 0) - return -nu * P_s * (C_R * A_over_m) * r_star / norm(r_star) + return -nu * P_s * (C_R * A_over_m) * r_star / norm_hf(_arr2tup_hf(r_star)) diff --git a/src/hapsira/core/propagation/vallado.py b/src/hapsira/core/propagation/vallado.py index c6b6f2914..c9a846641 100644 --- a/src/hapsira/core/propagation/vallado.py +++ b/src/hapsira/core/propagation/vallado.py @@ -1,7 +1,8 @@ from numba import njit as jit import numpy as np -from ..math.linalg import norm +from ..jit import _arr2tup_hf +from ..math.linalg import norm_hf from ..math.special import stumpff_c2 as c2, stumpff_c3 as c3 @@ -73,7 +74,7 @@ def vallado(k, r0, v0, tof, numiter): """ # Cache some results dot_r0v0 = r0 @ v0 - norm_r0 = norm(r0) + norm_r0 = norm_hf(_arr2tup_hf(r0)) sqrt_mu = k**0.5 alpha = -(v0 @ v0) / k + 2 / norm_r0 diff --git a/src/hapsira/core/spheroid_location.py b/src/hapsira/core/spheroid_location.py index 3f49bd38e..588f03b46 100644 --- a/src/hapsira/core/spheroid_location.py +++ b/src/hapsira/core/spheroid_location.py @@ -3,14 +3,15 @@ from numba import njit as jit import numpy as np -from .math.linalg import norm +from .jit import _arr2tup_hf +from .math.linalg import norm_hf @jit def cartesian_cords(a, c, lon, lat, h): """Calculates cartesian coordinates. - Parameters + Parametersnorm_hf ---------- a : float Semi-major axis @@ -66,7 +67,7 @@ def N(a, b, c, cartesian_cords): """ x, y, z = cartesian_cords N = np.array([2 * x / a**2, 2 * y / b**2, 2 * z / c**2]) - N /= norm(N) + N /= norm_hf(_arr2tup_hf(N)) return N @@ -82,7 +83,7 @@ def tangential_vecs(N): """ u = np.array([1.0, 0, 0]) u -= (u @ N) * N - u /= norm(u) + u /= norm_hf(_arr2tup_hf(u)) v = np.cross(N, u) return u, v @@ -125,7 +126,7 @@ def distance(cartesian_cords, px, py, pz): """ c = cartesian_cords u = np.array([px, py, pz]) - d = norm(c - u) + d = norm_hf(_arr2tup_hf(c - u)) return d diff --git a/src/hapsira/core/thrust/change_a_inc.py b/src/hapsira/core/thrust/change_a_inc.py index a97c66fb2..2fb5ead90 100644 --- a/src/hapsira/core/thrust/change_a_inc.py +++ b/src/hapsira/core/thrust/change_a_inc.py @@ -4,7 +4,8 @@ from hapsira.core.elements import circular_velocity -from ..math.linalg import norm +from ..jit import _arr2tup_hf +from ..math.linalg import norm_hf @jit @@ -101,8 +102,8 @@ def a_d(t0, u_, k): # Change sign of beta with the out-of-plane velocity beta_ = beta(t0, V_0, f, beta_0_) * np.sign(r[0] * (inc_f - inc_0)) - t_ = v / norm(v) - w_ = cross(r, v) / norm(cross(r, v)) + t_ = v / norm_hf(_arr2tup_hf(v)) + w_ = cross(r, v) / norm_hf(_arr2tup_hf(cross(r, v))) accel_v = f * (np.cos(beta_) * t_ + np.sin(beta_) * w_) return accel_v diff --git a/src/hapsira/core/thrust/change_argp.py b/src/hapsira/core/thrust/change_argp.py index 4fc8da570..6955161fb 100644 --- a/src/hapsira/core/thrust/change_argp.py +++ b/src/hapsira/core/thrust/change_argp.py @@ -4,7 +4,8 @@ from hapsira.core.elements import circular_velocity, rv2coe -from ..math.linalg import norm +from ..jit import _arr2tup_hf +from ..math.linalg import norm_hf @jit @@ -61,8 +62,8 @@ def a_d(t0, u_, k): alpha_ = nu - np.pi / 2 - r_ = r / norm(r) - w_ = cross(r, v) / norm(cross(r, v)) + r_ = r / norm_hf(_arr2tup_hf(r)) + w_ = cross(r, v) / norm_hf(_arr2tup_hf(cross(r, v))) s_ = cross(w_, r_) accel_v = f * (np.cos(alpha_) * s_ + np.sin(alpha_) * r_) return accel_v diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index 61fa04ea1..e6b27bff0 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -15,7 +15,8 @@ rv2coe, ) -from ..math.linalg import norm +from ..jit import _arr2tup_hf +from ..math.linalg import norm_hf @jit @@ -60,10 +61,10 @@ def change_ecc_inc(k, a, ecc_0, ecc_f, inc_0, inc_f, argp, r, v, f): e_vec = eccentricity_vector(k, r, v) ref_vec = e_vec / ecc_0 else: - ref_vec = r / norm(r) + ref_vec = r / norm_hf(_arr2tup_hf(r)) h_vec = cross(r, v) # Specific angular momentum vector - h_unit = h_vec / norm(h_vec) + h_unit = h_vec / norm_hf(_arr2tup_hf(h_vec)) thrust_unit = cross(h_unit, ref_vec) * np.sign(ecc_f - ecc_0) beta_0 = beta(ecc_0, ecc_f, inc_0, inc_f, argp) @@ -77,7 +78,9 @@ def a_d(t0, u_, k_): np.cos(nu) ) # The sign of ß reverses at minor axis crossings - w_ = (cross(r_, v_) / norm(cross(r_, v_))) * np.sign(inc_f - inc_0) + w_ = (cross(r_, v_) / norm_hf(_arr2tup_hf(cross(r_, v_)))) * np.sign( + inc_f - inc_0 + ) accel_v = f * (np.cos(beta_) * thrust_unit + np.sin(beta_) * w_) return accel_v diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index ab54e952c..3da687531 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -4,7 +4,7 @@ from astropy.coordinates import get_body_barycentric_posvel import numpy as np -from hapsira.core.math.linalg import norm +from hapsira.core.math.linalg import norm_vf from hapsira.core.events import ( eclipse_function as eclipse_function_fast, line_of_sight as line_of_sight_fast, @@ -71,7 +71,7 @@ def __init__(self, alt, R, terminal=True, direction=-1): def __call__(self, t, u, k): self._last_t = t - r_norm = norm(u[:3]) + r_norm = norm_vf(*u[:3]) return ( r_norm - self._R - self._alt @@ -121,7 +121,7 @@ def __init__(self, orbit, lat, terminal=False, direction=0): def __call__(self, t, u_, k): self._last_t = t - pos_on_body = (u_[:3] / norm(u_[:3])) * self._R + pos_on_body = (u_[:3] / norm_vf(*u_[:3])) * self._R _, lat_, _ = cartesian_to_ellipsoidal_fast(self._R, self._R_polar, *pos_on_body) return np.rad2deg(lat_) - self._lat @@ -273,7 +273,7 @@ def __init__(self, attractor, pos_coords, terminal=False, direction=0): def __call__(self, t, u_, k): self._last_t = t - if norm(u_[:3]) < self._R: + if norm_vf(*u_[:3]) < self._R: warn( "The norm of the position vector of the primary body is less than the radius of the attractor." ) diff --git a/src/hapsira/util.py b/src/hapsira/util.py index 5b56137d0..074e0bad5 100644 --- a/src/hapsira/util.py +++ b/src/hapsira/util.py @@ -4,7 +4,7 @@ from astropy.time import Time import numpy as np -from hapsira.core.math.linalg import norm as norm_fast +from hapsira.core.math.linalg import norm_vf from hapsira.core.util import alinspace as alinspace_fast @@ -26,7 +26,7 @@ def norm(vec, axis=None): result = norm_np(vec.value, axis=axis) else: - result = norm_fast(vec.value) + result = norm_vf(*vec.value) return result << vec.unit From eb1abf3b6eb03fe278eaa046575df0ad0296ca22 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 12:07:59 +0100 Subject: [PATCH 042/346] jit hpp21b --- src/hapsira/core/iod.py | 4 +-- src/hapsira/core/math/special.py | 44 ++++++++++++++++++++------------ tests/test_hyper.py | 4 +-- 3 files changed, 32 insertions(+), 20 deletions(-) diff --git a/src/hapsira/core/iod.py b/src/hapsira/core/iod.py index c3085d347..cab093bbc 100644 --- a/src/hapsira/core/iod.py +++ b/src/hapsira/core/iod.py @@ -4,7 +4,7 @@ from .jit import _arr2tup_hf from .math.linalg import norm_hf -from .math.special import hyp2f1b, stumpff_c2 as c2, stumpff_c3 as c3 +from .math.special import hyp2f1b_hf, stumpff_c2 as c2, stumpff_c3 as c3 @jit @@ -342,7 +342,7 @@ def _tof_equation_y(x, y, T0, ll, M): if M == 0 and np.sqrt(0.6) < x < np.sqrt(1.4): eta = y - ll * x S_1 = (1 - ll - x * eta) * 0.5 - Q = 4 / 3 * hyp2f1b(S_1) + Q = 4 / 3 * hyp2f1b_hf(S_1) T_ = (eta**3 * Q + 4 * ll * eta) * 0.5 else: psi = _compute_psi(x, y, ll) diff --git a/src/hapsira/core/math/special.py b/src/hapsira/core/math/special.py index 375cce97f..15f7f3afe 100644 --- a/src/hapsira/core/math/special.py +++ b/src/hapsira/core/math/special.py @@ -1,17 +1,20 @@ -from math import gamma +from math import gamma, inf from numba import njit as jit import numpy as np +from ..jit import hjit, vjit + __all__ = [ - "hyp2f1b", + "hyp2f1b_hf", + "hyp2f1b_vf", "stumpff_c2", "stumpff_c3", ] -@jit -def hyp2f1b(x): +@hjit("f(f)") +def hyp2f1b_hf(x): """Hypergeometric function 2F1(3, 1, 5/2, x), see [Battin]. .. todo:: @@ -24,18 +27,27 @@ def hyp2f1b(x): """ if x >= 1.0: - return np.inf - else: - res = 1.0 - term = 1.0 - ii = 0 - while True: - term = term * (3 + ii) * (1 + ii) / (5 / 2 + ii) * x / (ii + 1) - res_old = res - res += term - if res_old == res: - return res - ii += 1 + return inf + + res = 1.0 + term = 1.0 + ii = 0 + while True: + term = term * (3 + ii) * (1 + ii) / (5 / 2 + ii) * x / (ii + 1) + res_old = res + res += term + if res_old == res: + return res + ii += 1 + + +@vjit("f(f)") +def hyp2f1b_vf(x): + """ + Vectorized hyp2f1b + """ + + return hyp2f1b_hf(x) @jit diff --git a/tests/test_hyper.py b/tests/test_hyper.py index edeea790a..4baae59d3 100644 --- a/tests/test_hyper.py +++ b/tests/test_hyper.py @@ -3,12 +3,12 @@ import pytest from scipy import special -from hapsira.core.math.special import hyp2f1b as hyp2f1 +from hapsira.core.math.special import hyp2f1b_vf @pytest.mark.parametrize("x", np.linspace(0, 1, num=11)) def test_hyp2f1_battin_scalar(x): expected_res = special.hyp2f1(3, 1, 5 / 2, x) - res = hyp2f1(x) + res = hyp2f1b_vf(x) assert_allclose(res, expected_res) From ecb5ccf0a74c4ac3364241a40ed5cd60fd255353 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 12:19:25 +0100 Subject: [PATCH 043/346] jit stumpff_c2 --- src/hapsira/core/iod.py | 13 +++++--- src/hapsira/core/math/special.py | 43 ++++++++++++++++--------- src/hapsira/core/propagation/vallado.py | 4 +-- tests/test_stumpff.py | 8 ++--- 4 files changed, 41 insertions(+), 27 deletions(-) diff --git a/src/hapsira/core/iod.py b/src/hapsira/core/iod.py index cab093bbc..82b8b6448 100644 --- a/src/hapsira/core/iod.py +++ b/src/hapsira/core/iod.py @@ -4,7 +4,7 @@ from .jit import _arr2tup_hf from .math.linalg import norm_hf -from .math.special import hyp2f1b_hf, stumpff_c2 as c2, stumpff_c3 as c3 +from .math.special import hyp2f1b_hf, stumpff_c2_hf, stumpff_c3 as c3 @jit @@ -130,7 +130,7 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): count = 0 while count < numiter: - y = norm_r0_plus_norm_r + A * (psi * c3(psi) - 1) / c2(psi) ** 0.5 + y = norm_r0_plus_norm_r + A * (psi * c3(psi) - 1) / stumpff_c2_hf(psi) ** 0.5 if A > 0.0: # Readjust xi_low until y > 0.0 # Translated directly from Vallado @@ -139,11 +139,14 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): psi = ( 0.8 * (1.0 / c3(psi)) - * (1.0 - norm_r0_times_norm_r * np.sqrt(c2(psi)) / A) + * (1.0 - norm_r0_times_norm_r * np.sqrt(stumpff_c2_hf(psi)) / A) + ) + y = ( + norm_r0_plus_norm_r + + A * (psi * c3(psi) - 1) / stumpff_c2_hf(psi) ** 0.5 ) - y = norm_r0_plus_norm_r + A * (psi * c3(psi) - 1) / c2(psi) ** 0.5 - xi = np.sqrt(y / c2(psi)) + xi = np.sqrt(y / stumpff_c2_hf(psi)) tof_new = (xi**3 * c3(psi) + A * np.sqrt(y)) / np.sqrt(k) # Convergence check diff --git a/src/hapsira/core/math/special.py b/src/hapsira/core/math/special.py index 15f7f3afe..2082baeac 100644 --- a/src/hapsira/core/math/special.py +++ b/src/hapsira/core/math/special.py @@ -1,4 +1,4 @@ -from math import gamma, inf +from math import cos, cosh, gamma, inf, sqrt from numba import njit as jit import numpy as np @@ -8,7 +8,8 @@ __all__ = [ "hyp2f1b_hf", "hyp2f1b_vf", - "stumpff_c2", + "stumpff_c2_hf", + "stumpff_c2_vf", "stumpff_c3", ] @@ -50,8 +51,8 @@ def hyp2f1b_vf(x): return hyp2f1b_hf(x) -@jit -def stumpff_c2(psi): +@hjit("f(f)") +def stumpff_c2_hf(psi): r"""Second Stumpff function. For positive arguments: @@ -62,22 +63,32 @@ def stumpff_c2(psi): """ eps = 1.0 - if psi > eps: - res = (1 - np.cos(np.sqrt(psi))) / psi - elif psi < -eps: - res = (np.cosh(np.sqrt(-psi)) - 1) / (-psi) - else: - res = 1.0 / 2.0 - delta = (-psi) / gamma(2 + 2 + 1) - k = 1 - while res + delta != res: - res = res + delta - k += 1 - delta = (-psi) ** k / gamma(2 * k + 2 + 1) + if psi > eps: + return (1 - cos(sqrt(psi))) / psi + + if psi < -eps: + return (cosh(sqrt(-psi)) - 1) / (-psi) + + res = 1.0 / 2.0 + delta = (-psi) / gamma(2 + 2 + 1) + k = 1 + while res + delta != res: + res = res + delta + k += 1 + delta = (-psi) ** k / gamma(2 * k + 2 + 1) return res +@vjit("f(f)") +def stumpff_c2_vf(psi): + """ + Vectorized stumpff_c2 + """ + + return stumpff_c2_hf(psi) + + @jit def stumpff_c3(psi): r"""Third Stumpff function. diff --git a/src/hapsira/core/propagation/vallado.py b/src/hapsira/core/propagation/vallado.py index c9a846641..805ec81e2 100644 --- a/src/hapsira/core/propagation/vallado.py +++ b/src/hapsira/core/propagation/vallado.py @@ -3,7 +3,7 @@ from ..jit import _arr2tup_hf from ..math.linalg import norm_hf -from ..math.special import stumpff_c2 as c2, stumpff_c3 as c3 +from ..math.special import stumpff_c2_hf, stumpff_c3 as c3 @jit @@ -105,7 +105,7 @@ def vallado(k, r0, v0, tof, numiter): while count < numiter: xi = xi_new psi = xi * xi * alpha - c2_psi = c2(psi) + c2_psi = stumpff_c2_hf(psi) c3_psi = c3(psi) norm_r = ( xi * xi * c2_psi diff --git a/tests/test_stumpff.py b/tests/test_stumpff.py index 60ccd6830..b37d8fbff 100644 --- a/tests/test_stumpff.py +++ b/tests/test_stumpff.py @@ -1,7 +1,7 @@ from numpy import cos, cosh, sin, sinh from numpy.testing import assert_allclose -from hapsira.core.math.special import stumpff_c2 as c2, stumpff_c3 as c3 +from hapsira.core.math.special import stumpff_c2_vf, stumpff_c3 as c3 def test_stumpff_functions_near_zero(): @@ -9,7 +9,7 @@ def test_stumpff_functions_near_zero(): expected_c2 = (1 - cos(psi**0.5)) / psi expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5 - assert_allclose(c2(psi), expected_c2) + assert_allclose(stumpff_c2_vf(psi), expected_c2) assert_allclose(c3(psi), expected_c3) @@ -18,7 +18,7 @@ def test_stumpff_functions_above_zero(): expected_c2 = (1 - cos(psi**0.5)) / psi expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5 - assert_allclose(c2(psi), expected_c2, rtol=1e-10) + assert_allclose(stumpff_c2_vf(psi), expected_c2, rtol=1e-10) assert_allclose(c3(psi), expected_c3, rtol=1e-10) @@ -27,5 +27,5 @@ def test_stumpff_functions_under_zero(): expected_c2 = (cosh((-psi) ** 0.5) - 1) / (-psi) expected_c3 = (sinh((-psi) ** 0.5) - (-psi) ** 0.5) / (-psi) ** 1.5 - assert_allclose(c2(psi), expected_c2, rtol=1e-10) + assert_allclose(stumpff_c2_vf(psi), expected_c2, rtol=1e-10) assert_allclose(c3(psi), expected_c3, rtol=1e-10) From 1f5713e34d362b5cc2295bf298b667139d0261a8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 12:34:40 +0100 Subject: [PATCH 044/346] jit stumpff_c3 --- src/hapsira/core/iod.py | 13 +++++--- src/hapsira/core/math/special.py | 44 +++++++++++++++---------- src/hapsira/core/propagation/vallado.py | 4 +-- tests/test_stumpff.py | 8 ++--- 4 files changed, 40 insertions(+), 29 deletions(-) diff --git a/src/hapsira/core/iod.py b/src/hapsira/core/iod.py index 82b8b6448..364969173 100644 --- a/src/hapsira/core/iod.py +++ b/src/hapsira/core/iod.py @@ -4,7 +4,7 @@ from .jit import _arr2tup_hf from .math.linalg import norm_hf -from .math.special import hyp2f1b_hf, stumpff_c2_hf, stumpff_c3 as c3 +from .math.special import hyp2f1b_hf, stumpff_c2_hf, stumpff_c3_hf @jit @@ -130,7 +130,10 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): count = 0 while count < numiter: - y = norm_r0_plus_norm_r + A * (psi * c3(psi) - 1) / stumpff_c2_hf(psi) ** 0.5 + y = ( + norm_r0_plus_norm_r + + A * (psi * stumpff_c3_hf(psi) - 1) / stumpff_c2_hf(psi) ** 0.5 + ) if A > 0.0: # Readjust xi_low until y > 0.0 # Translated directly from Vallado @@ -138,16 +141,16 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): psi_low = psi psi = ( 0.8 - * (1.0 / c3(psi)) + * (1.0 / stumpff_c3_hf(psi)) * (1.0 - norm_r0_times_norm_r * np.sqrt(stumpff_c2_hf(psi)) / A) ) y = ( norm_r0_plus_norm_r - + A * (psi * c3(psi) - 1) / stumpff_c2_hf(psi) ** 0.5 + + A * (psi * stumpff_c3_hf(psi) - 1) / stumpff_c2_hf(psi) ** 0.5 ) xi = np.sqrt(y / stumpff_c2_hf(psi)) - tof_new = (xi**3 * c3(psi) + A * np.sqrt(y)) / np.sqrt(k) + tof_new = (xi**3 * stumpff_c3_hf(psi) + A * np.sqrt(y)) / np.sqrt(k) # Convergence check if np.abs((tof_new - tof) / tof) < rtol: diff --git a/src/hapsira/core/math/special.py b/src/hapsira/core/math/special.py index 2082baeac..34f142181 100644 --- a/src/hapsira/core/math/special.py +++ b/src/hapsira/core/math/special.py @@ -1,7 +1,4 @@ -from math import cos, cosh, gamma, inf, sqrt - -from numba import njit as jit -import numpy as np +from math import cos, cosh, gamma, inf, sin, sinh, sqrt from ..jit import hjit, vjit @@ -10,7 +7,8 @@ "hyp2f1b_vf", "stumpff_c2_hf", "stumpff_c2_vf", - "stumpff_c3", + "stumpff_c3_hf", + "stumpff_c3_vf", ] @@ -89,8 +87,8 @@ def stumpff_c2_vf(psi): return stumpff_c2_hf(psi) -@jit -def stumpff_c3(psi): +@hjit("f(f)") +def stumpff_c3_hf(psi): r"""Third Stumpff function. For positive arguments: @@ -101,17 +99,27 @@ def stumpff_c3(psi): """ eps = 1.0 + if psi > eps: - res = (np.sqrt(psi) - np.sin(np.sqrt(psi))) / (psi * np.sqrt(psi)) - elif psi < -eps: - res = (np.sinh(np.sqrt(-psi)) - np.sqrt(-psi)) / (-psi * np.sqrt(-psi)) - else: - res = 1.0 / 6.0 - delta = (-psi) / gamma(2 + 3 + 1) - k = 1 - while res + delta != res: - res = res + delta - k += 1 - delta = (-psi) ** k / gamma(2 * k + 3 + 1) + return (sqrt(psi) - sin(sqrt(psi))) / (psi * sqrt(psi)) + + if psi < -eps: + return (sinh(sqrt(-psi)) - sqrt(-psi)) / (-psi * sqrt(-psi)) + res = 1.0 / 6.0 + delta = (-psi) / gamma(2 + 3 + 1) + k = 1 + while res + delta != res: + res = res + delta + k += 1 + delta = (-psi) ** k / gamma(2 * k + 3 + 1) return res + + +@vjit("f(f)") +def stumpff_c3_vf(psi): + """ + Vectorized stumpff_c3 + """ + + return stumpff_c3_hf(psi) diff --git a/src/hapsira/core/propagation/vallado.py b/src/hapsira/core/propagation/vallado.py index 805ec81e2..f0b33dbf8 100644 --- a/src/hapsira/core/propagation/vallado.py +++ b/src/hapsira/core/propagation/vallado.py @@ -3,7 +3,7 @@ from ..jit import _arr2tup_hf from ..math.linalg import norm_hf -from ..math.special import stumpff_c2_hf, stumpff_c3 as c3 +from ..math.special import stumpff_c2_hf, stumpff_c3_hf @jit @@ -106,7 +106,7 @@ def vallado(k, r0, v0, tof, numiter): xi = xi_new psi = xi * xi * alpha c2_psi = stumpff_c2_hf(psi) - c3_psi = c3(psi) + c3_psi = stumpff_c3_hf(psi) norm_r = ( xi * xi * c2_psi + dot_r0v0 / sqrt_mu * xi * (1 - psi * c3_psi) diff --git a/tests/test_stumpff.py b/tests/test_stumpff.py index b37d8fbff..621867884 100644 --- a/tests/test_stumpff.py +++ b/tests/test_stumpff.py @@ -1,7 +1,7 @@ from numpy import cos, cosh, sin, sinh from numpy.testing import assert_allclose -from hapsira.core.math.special import stumpff_c2_vf, stumpff_c3 as c3 +from hapsira.core.math.special import stumpff_c2_vf, stumpff_c3_vf def test_stumpff_functions_near_zero(): @@ -10,7 +10,7 @@ def test_stumpff_functions_near_zero(): expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5 assert_allclose(stumpff_c2_vf(psi), expected_c2) - assert_allclose(c3(psi), expected_c3) + assert_allclose(stumpff_c3_vf(psi), expected_c3) def test_stumpff_functions_above_zero(): @@ -19,7 +19,7 @@ def test_stumpff_functions_above_zero(): expected_c3 = (psi**0.5 - sin(psi**0.5)) / psi**1.5 assert_allclose(stumpff_c2_vf(psi), expected_c2, rtol=1e-10) - assert_allclose(c3(psi), expected_c3, rtol=1e-10) + assert_allclose(stumpff_c3_vf(psi), expected_c3, rtol=1e-10) def test_stumpff_functions_under_zero(): @@ -28,4 +28,4 @@ def test_stumpff_functions_under_zero(): expected_c3 = (sinh((-psi) ** 0.5) - (-psi) ** 0.5) / (-psi) ** 1.5 assert_allclose(stumpff_c2_vf(psi), expected_c2, rtol=1e-10) - assert_allclose(c3(psi), expected_c3, rtol=1e-10) + assert_allclose(stumpff_c3_vf(psi), expected_c3, rtol=1e-10) From 36c66ce7146e01b7b4488939b157fb57706a2af7 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 12:39:09 +0100 Subject: [PATCH 045/346] exports --- src/hapsira/core/elements.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 49d228fd2..1c8c819b1 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -15,6 +15,20 @@ from .math.linalg import norm_hf +__all__ = [ + "eccentricity_vector", + "circular_velocity", + "rv_pqw", + "coe_rotation_matrix", + "coe2rv", + "coe2rv_many", + "coe2mee", + "rv2coe", + "mee2coe", + "mee2rv", +] + + @jit def eccentricity_vector(k, r, v): r"""Eccentricity vector. From f75877f42ab7bc97d74266568d0b987bab2619f7 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 13:30:40 +0100 Subject: [PATCH 046/346] jit eccentricity_vector --- src/hapsira/core/elements.py | 35 +++++++++++++++++------ src/hapsira/core/jit.py | 1 + src/hapsira/core/math/linalg.py | 18 ++++++++++++ src/hapsira/core/thrust/change_ecc_inc.py | 5 ++-- src/hapsira/twobody/elements.py | 11 +++---- 5 files changed, 53 insertions(+), 17 deletions(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 1c8c819b1..15e893421 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -11,12 +11,20 @@ from hapsira.core.angles import E_to_nu_hf, F_to_nu_hf from hapsira.core.util import rotation_matrix -from .jit import _arr2tup_hf -from .math.linalg import norm_hf +from .jit import _arr2tup_hf, hjit, gjit +from .math.linalg import ( + div_Vs_hf, + matmul_VV_hf, + mul_Vs_hf, + norm_hf, + sub_VV_hf, +) __all__ = [ - "eccentricity_vector", + "eccentricity_vector_hf", + "eccentricity_vector_gf", + # TODO "circular_velocity", "rv_pqw", "coe_rotation_matrix", @@ -29,8 +37,8 @@ ] -@jit -def eccentricity_vector(k, r, v): +@hjit("V(f,V,V)") +def eccentricity_vector_hf(k, r, v): r"""Eccentricity vector. .. math:: @@ -43,12 +51,23 @@ def eccentricity_vector(k, r, v): ---------- k : float Standard gravitational parameter (km^3 / s^2). - r : numpy.ndarray + r : tuple[float,float,float] Position vector (km) - v : numpy.ndarray + v : tuple[float,float,float] Velocity vector (km / s) """ - return ((v @ v - k / norm_hf(_arr2tup_hf(r))) * r - (r @ v) * v) / k + a = matmul_VV_hf(v, v) - k / norm_hf(r) + b = matmul_VV_hf(r, v) + return div_Vs_hf(sub_VV_hf(mul_Vs_hf(r, a), mul_Vs_hf(v, b)), k) + + +@gjit("void(f,f[:],f[:],f[:])", "(),(n),(n)->(n)") +def eccentricity_vector_gf(k, r, v, e): + """ + Vectorized eccentricity_vector + """ + + e[0], e[1], e[2] = eccentricity_vector_hf(k, _arr2tup_hf(r), _arr2tup_hf(v)) @jit diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index b4a2aea3c..bca13a196 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -14,6 +14,7 @@ "vjit", "gjit", "sjit", + "_arr2tup_hf", ] diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 28404cd3d..309996d78 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -3,12 +3,25 @@ from ..jit import hjit, vjit __all__ = [ + "div_Vs_hf", "matmul_VV_hf", + "mul_Vs_hf", "norm_hf", "norm_vf", + "sub_VV_hf", ] +@hjit("V(V,f)") +def div_Vs_hf(v, s): + return v[0] / s, v[1] / s, v[2] / s + + +@hjit("V(V,f)") +def mul_Vs_hf(v, s): + return v[0] * s, v[1] * s, v[2] * s + + @hjit("f(V,V)") def matmul_VV_hf(a, b): return a[0] * b[0] + a[1] * b[1] + a[2] * b[2] @@ -23,3 +36,8 @@ def norm_hf(a): def norm_vf(a, b, c): # TODO add axis setting in some way for util.norm? return norm_hf((a, b, c)) + + +@hjit("V(V,V)") +def sub_VV_hf(va, vb): + return va[0] - vb[0], va[1] - vb[1], va[2] - vb[2] diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index e6b27bff0..4526baa2d 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -11,7 +11,7 @@ from hapsira.core.elements import ( circular_velocity, - eccentricity_vector, + eccentricity_vector_gf, rv2coe, ) @@ -58,7 +58,8 @@ def delta_t(delta_v, f): def change_ecc_inc(k, a, ecc_0, ecc_f, inc_0, inc_f, argp, r, v, f): # We fix the inertial direction at the beginning if ecc_0 > 0.001: # Arbitrary tolerance - e_vec = eccentricity_vector(k, r, v) + e_vec = np.zeros_like(r) + eccentricity_vector_gf(k, r, v, e_vec) ref_vec = e_vec / ecc_0 else: ref_vec = r / norm_hf(_arr2tup_hf(r)) diff --git a/src/hapsira/twobody/elements.py b/src/hapsira/twobody/elements.py index 7b14f08d3..3f6f55711 100644 --- a/src/hapsira/twobody/elements.py +++ b/src/hapsira/twobody/elements.py @@ -5,7 +5,7 @@ circular_velocity as circular_velocity_fast, coe2rv as coe2rv_fast, coe2rv_many as coe2rv_many_fast, - eccentricity_vector as eccentricity_vector_fast, + eccentricity_vector_gf, ) from hapsira.core.propagation.farnocchia import ( delta_t_from_nu as delta_t_from_nu_fast, @@ -43,12 +43,9 @@ def energy(k, r, v): @u.quantity_input(k=u_km3s2, r=u.km, v=u_kms) def eccentricity_vector(k, r, v): """Eccentricity vector.""" - return ( - eccentricity_vector_fast( - k.to_value(u_km3s2), r.to_value(u.km), v.to_value(u_kms) - ) - * u.one - ) + e = np.zeros(r.shape, dtype=r.dtype) + eccentricity_vector_gf(k.to_value(u_km3s2), r.to_value(u.km), v.to_value(u_kms), e) + return e << u.one @u.quantity_input(nu=u.rad, ecc=u.one, k=u_km3s2, r_p=u.km) From 98e7cdf234251462e0a9c1ab2bdfac066e66afad Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 13:53:07 +0100 Subject: [PATCH 047/346] rename helper func, keep it --- src/hapsira/core/elements.py | 20 ++++++++++---------- src/hapsira/core/events.py | 8 ++++---- src/hapsira/core/flybys.py | 8 ++++---- src/hapsira/core/iod.py | 14 +++++++------- src/hapsira/core/jit.py | 6 +++--- src/hapsira/core/maneuver.py | 16 ++++++++-------- src/hapsira/core/perturbations.py | 20 ++++++++++---------- src/hapsira/core/propagation/vallado.py | 4 ++-- src/hapsira/core/spheroid_location.py | 8 ++++---- src/hapsira/core/thrust/change_a_inc.py | 6 +++--- src/hapsira/core/thrust/change_argp.py | 6 +++--- src/hapsira/core/thrust/change_ecc_inc.py | 8 ++++---- 12 files changed, 62 insertions(+), 62 deletions(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 15e893421..137e49593 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -11,7 +11,7 @@ from hapsira.core.angles import E_to_nu_hf, F_to_nu_hf from hapsira.core.util import rotation_matrix -from .jit import _arr2tup_hf, hjit, gjit +from .jit import array_to_V_hf, hjit, gjit from .math.linalg import ( div_Vs_hf, matmul_VV_hf, @@ -67,7 +67,7 @@ def eccentricity_vector_gf(k, r, v, e): Vectorized eccentricity_vector """ - e[0], e[1], e[2] = eccentricity_vector_hf(k, _arr2tup_hf(r), _arr2tup_hf(v)) + e[0], e[1], e[2] = eccentricity_vector_hf(k, array_to_V_hf(r), array_to_V_hf(v)) @jit @@ -415,10 +415,10 @@ def rv2coe(k, r, v, tol=1e-8): """ h = cross(r, v) n = cross([0, 0, 1], h) - e = ((v @ v - k / norm_hf(_arr2tup_hf(r))) * r - (r @ v) * v) / k - ecc = norm_hf(_arr2tup_hf(e)) + e = ((v @ v - k / norm_hf(array_to_V_hf(r))) * r - (r @ v) * v) / k + ecc = norm_hf(array_to_V_hf(e)) p = (h @ h) / k - inc = np.arccos(h[2] / norm_hf(_arr2tup_hf(h))) + inc = np.arccos(h[2] / norm_hf(array_to_V_hf(h))) circular = ecc < tol equatorial = abs(inc) < tol @@ -426,12 +426,12 @@ def rv2coe(k, r, v, tol=1e-8): if equatorial and not circular: raan = 0 argp = np.arctan2(e[1], e[0]) % (2 * np.pi) # Longitude of periapsis - nu = np.arctan2((h @ cross(e, r)) / norm_hf(_arr2tup_hf(h)), r @ e) + nu = np.arctan2((h @ cross(e, r)) / norm_hf(array_to_V_hf(h)), r @ e) elif not equatorial and circular: raan = np.arctan2(n[1], n[0]) % (2 * np.pi) argp = 0 # Argument of latitude - nu = np.arctan2((r @ cross(h, n)) / norm_hf(_arr2tup_hf(h)), r @ n) + nu = np.arctan2((r @ cross(h, n)) / norm_hf(array_to_V_hf(h)), r @ n) elif equatorial and circular: raan = 0 argp = 0 @@ -441,16 +441,16 @@ def rv2coe(k, r, v, tol=1e-8): ka = k * a if a > 0: e_se = (r @ v) / sqrt(ka) - e_ce = norm_hf(_arr2tup_hf(r)) * (v @ v) / k - 1 + e_ce = norm_hf(array_to_V_hf(r)) * (v @ v) / k - 1 nu = E_to_nu_hf(np.arctan2(e_se, e_ce), ecc) else: e_sh = (r @ v) / sqrt(-ka) - e_ch = norm_hf(_arr2tup_hf(r)) * (norm_hf(_arr2tup_hf(v)) ** 2) / k - 1 + e_ch = norm_hf(array_to_V_hf(r)) * (norm_hf(array_to_V_hf(v)) ** 2) / k - 1 nu = F_to_nu_hf(np.log((e_ch + e_sh) / (e_ch - e_sh)) / 2, ecc) raan = np.arctan2(n[1], n[0]) % (2 * np.pi) px = r @ n - py = (r @ cross(h, n)) / norm_hf(_arr2tup_hf(h)) + py = (r @ cross(h, n)) / norm_hf(array_to_V_hf(h)) argp = (np.arctan2(py, px) - nu) % (2 * np.pi) nu = (nu + np.pi) % (2 * np.pi) - np.pi diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index c5b5a9290..b661de935 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -4,7 +4,7 @@ from hapsira.core.elements import coe_rotation_matrix, rv2coe from hapsira.core.util import planetocentric_to_AltAz -from .jit import _arr2tup_hf +from .jit import array_to_V_hf from .math.linalg import norm_hf @@ -42,7 +42,7 @@ def eclipse_function(k, u_, r_sec, R_sec, R_primary, umbra=True): # Make arrays contiguous for faster dot product with numba. P_, Q_ = np.ascontiguousarray(PQW[:, 0]), np.ascontiguousarray(PQW[:, 1]) - r_sec_norm = norm_hf(_arr2tup_hf(r_sec)) + r_sec_norm = norm_hf(array_to_V_hf(r_sec)) beta = (P_ @ r_sec) / r_sec_norm zeta = (Q_ @ r_sec) / r_sec_norm @@ -79,8 +79,8 @@ def line_of_sight(r1, r2, R): located by r1 and r2, else negative. """ - r1_norm = norm_hf(_arr2tup_hf(r1)) - r2_norm = norm_hf(_arr2tup_hf(r2)) + r1_norm = norm_hf(array_to_V_hf(r1)) + r2_norm = norm_hf(array_to_V_hf(r2)) theta = np.arccos((r1 @ r2) / r1_norm / r2_norm) theta_1 = np.arccos(R / r1_norm) diff --git a/src/hapsira/core/flybys.py b/src/hapsira/core/flybys.py index 9a90fffff..18064ff0a 100644 --- a/src/hapsira/core/flybys.py +++ b/src/hapsira/core/flybys.py @@ -4,7 +4,7 @@ import numpy as np from numpy import cross -from .jit import _arr2tup_hf +from .jit import array_to_V_hf from .math.linalg import norm_hf @@ -34,7 +34,7 @@ def compute_flyby(v_spacecraft, v_body, k, r_p, theta): """ v_inf_1 = v_spacecraft - v_body # Hyperbolic excess velocity - v_inf = norm_hf(_arr2tup_hf(v_inf_1)) + v_inf = norm_hf(array_to_V_hf(v_inf_1)) ecc = 1 + r_p * v_inf**2 / k # Eccentricity of the entry hyperbola delta = 2 * np.arcsin(1 / ecc) # Turn angle @@ -49,7 +49,7 @@ def compute_flyby(v_spacecraft, v_body, k, r_p, theta): S_vec = v_inf_1 / v_inf c_vec = np.array([0, 0, 1]) T_vec = cross(S_vec, c_vec) - T_vec = T_vec / norm_hf(_arr2tup_hf(T_vec)) + T_vec = T_vec / norm_hf(array_to_V_hf(T_vec)) R_vec = cross(S_vec, T_vec) # This vector defines the B-Plane @@ -63,7 +63,7 @@ def compute_flyby(v_spacecraft, v_body, k, r_p, theta): # And now we rotate the outbound hyperbolic excess velocity # u_vec = v_inf_1 / norm(v_inf) = S_vec v_vec = cross(rot_v, v_inf_1) - v_vec = v_vec / norm_hf(_arr2tup_hf(v_vec)) + v_vec = v_vec / norm_hf(array_to_V_hf(v_vec)) v_inf_2 = v_inf * (np.cos(delta) * S_vec + np.sin(delta) * v_vec) diff --git a/src/hapsira/core/iod.py b/src/hapsira/core/iod.py index 364969173..3b07140ec 100644 --- a/src/hapsira/core/iod.py +++ b/src/hapsira/core/iod.py @@ -2,7 +2,7 @@ import numpy as np from numpy import cross, pi -from .jit import _arr2tup_hf +from .jit import array_to_V_hf from .math.linalg import norm_hf from .math.special import hyp2f1b_hf, stumpff_c2_hf, stumpff_c3_hf @@ -111,8 +111,8 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): t_m = 1 if prograde else -1 - norm_r0 = norm_hf(_arr2tup_hf(r0)) - norm_r = norm_hf(_arr2tup_hf(r)) + norm_r0 = norm_hf(array_to_V_hf(r0)) + norm_r = norm_hf(array_to_V_hf(r)) norm_r0_times_norm_r = norm_r0 * norm_r norm_r0_plus_norm_r = norm_r0 + norm_r @@ -221,9 +221,9 @@ def izzo(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol): # Chord c = r2 - r1 c_norm, r1_norm, r2_norm = ( - norm_hf(_arr2tup_hf(c)), - norm_hf(_arr2tup_hf(r1)), - norm_hf(_arr2tup_hf(r2)), + norm_hf(array_to_V_hf(c)), + norm_hf(array_to_V_hf(r1)), + norm_hf(array_to_V_hf(r2)), ) # Semiperimeter @@ -232,7 +232,7 @@ def izzo(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol): # Versors i_r1, i_r2 = r1 / r1_norm, r2 / r2_norm i_h = cross(i_r1, i_r2) - i_h = i_h / norm_hf(_arr2tup_hf(i_h)) # Fixed from paper + i_h = i_h / norm_hf(array_to_V_hf(i_h)) # Fixed from paper # Geometry of the problem ll = np.sqrt(1 - min(1.0, c_norm / s)) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index bca13a196..a98d4ffd1 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -14,7 +14,7 @@ "vjit", "gjit", "sjit", - "_arr2tup_hf", + "array_to_V_hf", ] @@ -252,6 +252,6 @@ def wrapper(inner_func: Callable) -> Callable: return wrapper -@hjit("V(f[:])") # TODO remove, only for code refactor -def _arr2tup_hf(x): +@hjit("V(f[:])") +def array_to_V_hf(x): return x[0], x[1], x[2] diff --git a/src/hapsira/core/maneuver.py b/src/hapsira/core/maneuver.py index 2aafc4e11..8ea963488 100644 --- a/src/hapsira/core/maneuver.py +++ b/src/hapsira/core/maneuver.py @@ -6,7 +6,7 @@ from hapsira.core.elements import coe_rotation_matrix, rv2coe, rv_pqw -from .jit import _arr2tup_hf +from .jit import array_to_V_hf from .math.linalg import norm_hf @@ -43,13 +43,13 @@ def hohmann(k, rv, r_f): """ _, ecc, inc, raan, argp, nu = rv2coe(k, *rv) - h_i = norm_hf(_arr2tup_hf(cross(*rv))) + h_i = norm_hf(array_to_V_hf(cross(*rv))) p_i = h_i**2 / k r_i, v_i = rv_pqw(k, p_i, ecc, nu) - r_i = norm_hf(_arr2tup_hf(r_i)) - v_i = norm_hf(_arr2tup_hf(v_i)) + r_i = norm_hf(array_to_V_hf(r_i)) + v_i = norm_hf(array_to_V_hf(v_i)) a_trans = (r_i + r_f) / 2 dv_a = np.sqrt(2 * k / r_i - k / a_trans) - v_i @@ -113,13 +113,13 @@ def bielliptic(k, r_b, r_f, rv): """ _, ecc, inc, raan, argp, nu = rv2coe(k, *rv) - h_i = norm_hf(_arr2tup_hf(cross(*rv))) + h_i = norm_hf(array_to_V_hf(cross(*rv))) p_i = h_i**2 / k r_i, v_i = rv_pqw(k, p_i, ecc, nu) - r_i = norm_hf(_arr2tup_hf(r_i)) - v_i = norm_hf(_arr2tup_hf(v_i)) + r_i = norm_hf(array_to_V_hf(r_i)) + v_i = norm_hf(array_to_V_hf(v_i)) a_trans1 = (r_i + r_b) / 2 a_trans2 = (r_b + r_f) / 2 @@ -192,6 +192,6 @@ def correct_pericenter(k, R, J2, max_delta_r, v, a, inc, ecc): delta_t = abs(delta_w / dw) delta_v = 0.5 * n * a * ecc * abs(delta_w) - vf_ = v / norm_hf(_arr2tup_hf(v)) * delta_v + vf_ = v / norm_hf(array_to_V_hf(v)) * delta_v return delta_t, vf_ diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index df621ef42..6006e5103 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -3,7 +3,7 @@ from hapsira.core.events import line_of_sight as line_of_sight_fast -from .jit import _arr2tup_hf +from .jit import array_to_V_hf from .math.linalg import norm_hf @@ -37,7 +37,7 @@ def J2_perturbation(t0, state, k, J2, R): """ r_vec = state[:3] - r = norm_hf(_arr2tup_hf(r_vec)) + r = norm_hf(array_to_V_hf(r_vec)) factor = (3.0 / 2.0) * k * J2 * (R**2) / (r**5) @@ -72,7 +72,7 @@ def J3_perturbation(t0, state, k, J3, R): """ r_vec = state[:3] - r = norm_hf(_arr2tup_hf(r_vec)) + r = norm_hf(array_to_V_hf(r_vec)) factor = (1.0 / 2.0) * k * J3 * (R**3) / (r**5) cos_phi = r_vec[2] / r @@ -120,10 +120,10 @@ def atmospheric_drag_exponential(t0, state, k, R, C_D, A_over_m, H0, rho0): the atmospheric density model is rho(H) = rho0 x exp(-H / H0) """ - H = norm_hf(_arr2tup_hf(state[:3])) + H = norm_hf(array_to_V_hf(state[:3])) v_vec = state[3:] - v = norm_hf(_arr2tup_hf(v_vec)) + v = norm_hf(array_to_V_hf(v_vec)) B = C_D * A_over_m rho = rho0 * np.exp(-(H - R) / H0) @@ -162,7 +162,7 @@ def atmospheric_drag(t0, state, k, C_D, A_over_m, rho): """ v_vec = state[3:] - v = norm_hf(_arr2tup_hf(v_vec)) + v = norm_hf(array_to_V_hf(v_vec)) B = C_D * A_over_m return -(1.0 / 2.0) * rho * B * v * v_vec @@ -198,8 +198,8 @@ def third_body(t0, state, k, k_third, perturbation_body): body_r = perturbation_body(t0) delta_r = body_r - state[:3] return ( - k_third * delta_r / norm_hf(_arr2tup_hf(delta_r)) ** 3 - - k_third * body_r / norm_hf(_arr2tup_hf(body_r)) ** 3 + k_third * delta_r / norm_hf(array_to_V_hf(delta_r)) ** 3 + - k_third * body_r / norm_hf(array_to_V_hf(body_r)) ** 3 ) @@ -238,7 +238,7 @@ def radiation_pressure(t0, state, k, R, C_R, A_over_m, Wdivc_s, star): """ r_star = star(t0) r_sat = state[:3] - P_s = Wdivc_s / (norm_hf(_arr2tup_hf(r_star)) ** 2) + P_s = Wdivc_s / (norm_hf(array_to_V_hf(r_star)) ** 2) nu = float(line_of_sight_fast(r_sat, r_star, R) > 0) - return -nu * P_s * (C_R * A_over_m) * r_star / norm_hf(_arr2tup_hf(r_star)) + return -nu * P_s * (C_R * A_over_m) * r_star / norm_hf(array_to_V_hf(r_star)) diff --git a/src/hapsira/core/propagation/vallado.py b/src/hapsira/core/propagation/vallado.py index f0b33dbf8..ce94dcf77 100644 --- a/src/hapsira/core/propagation/vallado.py +++ b/src/hapsira/core/propagation/vallado.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from ..jit import _arr2tup_hf +from ..jit import array_to_V_hf from ..math.linalg import norm_hf from ..math.special import stumpff_c2_hf, stumpff_c3_hf @@ -74,7 +74,7 @@ def vallado(k, r0, v0, tof, numiter): """ # Cache some results dot_r0v0 = r0 @ v0 - norm_r0 = norm_hf(_arr2tup_hf(r0)) + norm_r0 = norm_hf(array_to_V_hf(r0)) sqrt_mu = k**0.5 alpha = -(v0 @ v0) / k + 2 / norm_r0 diff --git a/src/hapsira/core/spheroid_location.py b/src/hapsira/core/spheroid_location.py index 588f03b46..86fe567fc 100644 --- a/src/hapsira/core/spheroid_location.py +++ b/src/hapsira/core/spheroid_location.py @@ -3,7 +3,7 @@ from numba import njit as jit import numpy as np -from .jit import _arr2tup_hf +from .jit import array_to_V_hf from .math.linalg import norm_hf @@ -67,7 +67,7 @@ def N(a, b, c, cartesian_cords): """ x, y, z = cartesian_cords N = np.array([2 * x / a**2, 2 * y / b**2, 2 * z / c**2]) - N /= norm_hf(_arr2tup_hf(N)) + N /= norm_hf(array_to_V_hf(N)) return N @@ -83,7 +83,7 @@ def tangential_vecs(N): """ u = np.array([1.0, 0, 0]) u -= (u @ N) * N - u /= norm_hf(_arr2tup_hf(u)) + u /= norm_hf(array_to_V_hf(u)) v = np.cross(N, u) return u, v @@ -126,7 +126,7 @@ def distance(cartesian_cords, px, py, pz): """ c = cartesian_cords u = np.array([px, py, pz]) - d = norm_hf(_arr2tup_hf(c - u)) + d = norm_hf(array_to_V_hf(c - u)) return d diff --git a/src/hapsira/core/thrust/change_a_inc.py b/src/hapsira/core/thrust/change_a_inc.py index 2fb5ead90..351aefab9 100644 --- a/src/hapsira/core/thrust/change_a_inc.py +++ b/src/hapsira/core/thrust/change_a_inc.py @@ -4,7 +4,7 @@ from hapsira.core.elements import circular_velocity -from ..jit import _arr2tup_hf +from ..jit import array_to_V_hf from ..math.linalg import norm_hf @@ -102,8 +102,8 @@ def a_d(t0, u_, k): # Change sign of beta with the out-of-plane velocity beta_ = beta(t0, V_0, f, beta_0_) * np.sign(r[0] * (inc_f - inc_0)) - t_ = v / norm_hf(_arr2tup_hf(v)) - w_ = cross(r, v) / norm_hf(_arr2tup_hf(cross(r, v))) + t_ = v / norm_hf(array_to_V_hf(v)) + w_ = cross(r, v) / norm_hf(array_to_V_hf(cross(r, v))) accel_v = f * (np.cos(beta_) * t_ + np.sin(beta_) * w_) return accel_v diff --git a/src/hapsira/core/thrust/change_argp.py b/src/hapsira/core/thrust/change_argp.py index 6955161fb..a66f59650 100644 --- a/src/hapsira/core/thrust/change_argp.py +++ b/src/hapsira/core/thrust/change_argp.py @@ -4,7 +4,7 @@ from hapsira.core.elements import circular_velocity, rv2coe -from ..jit import _arr2tup_hf +from ..jit import array_to_V_hf from ..math.linalg import norm_hf @@ -62,8 +62,8 @@ def a_d(t0, u_, k): alpha_ = nu - np.pi / 2 - r_ = r / norm_hf(_arr2tup_hf(r)) - w_ = cross(r, v) / norm_hf(_arr2tup_hf(cross(r, v))) + r_ = r / norm_hf(array_to_V_hf(r)) + w_ = cross(r, v) / norm_hf(array_to_V_hf(cross(r, v))) s_ = cross(w_, r_) accel_v = f * (np.cos(alpha_) * s_ + np.sin(alpha_) * r_) return accel_v diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index 4526baa2d..8b3f4be61 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -15,7 +15,7 @@ rv2coe, ) -from ..jit import _arr2tup_hf +from ..jit import array_to_V_hf from ..math.linalg import norm_hf @@ -62,10 +62,10 @@ def change_ecc_inc(k, a, ecc_0, ecc_f, inc_0, inc_f, argp, r, v, f): eccentricity_vector_gf(k, r, v, e_vec) ref_vec = e_vec / ecc_0 else: - ref_vec = r / norm_hf(_arr2tup_hf(r)) + ref_vec = r / norm_hf(array_to_V_hf(r)) h_vec = cross(r, v) # Specific angular momentum vector - h_unit = h_vec / norm_hf(_arr2tup_hf(h_vec)) + h_unit = h_vec / norm_hf(array_to_V_hf(h_vec)) thrust_unit = cross(h_unit, ref_vec) * np.sign(ecc_f - ecc_0) beta_0 = beta(ecc_0, ecc_f, inc_0, inc_f, argp) @@ -79,7 +79,7 @@ def a_d(t0, u_, k_): np.cos(nu) ) # The sign of ß reverses at minor axis crossings - w_ = (cross(r_, v_) / norm_hf(_arr2tup_hf(cross(r_, v_)))) * np.sign( + w_ = (cross(r_, v_) / norm_hf(array_to_V_hf(cross(r_, v_)))) * np.sign( inc_f - inc_0 ) accel_v = f * (np.cos(beta_) * thrust_unit + np.sin(beta_) * w_) From 32768c479319330936003601c971809d457ec69e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 14:33:38 +0100 Subject: [PATCH 048/346] jit circular_velocity --- src/hapsira/core/elements.py | 23 ++++++++++++++----- src/hapsira/core/thrust/change_a_inc.py | 6 ++--- src/hapsira/core/thrust/change_argp.py | 4 ++-- src/hapsira/core/thrust/change_ecc_inc.py | 4 ++-- .../core/thrust/change_ecc_quasioptimal.py | 4 ++-- src/hapsira/twobody/elements.py | 4 ++-- 6 files changed, 28 insertions(+), 17 deletions(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 137e49593..1c2879dfb 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -2,16 +2,17 @@ convert between different elements that define the orbit of a body. """ +from math import sqrt import sys from numba import njit as jit, prange import numpy as np -from numpy import cos, cross, sin, sqrt +from numpy import cos, cross, sin from hapsira.core.angles import E_to_nu_hf, F_to_nu_hf from hapsira.core.util import rotation_matrix -from .jit import array_to_V_hf, hjit, gjit +from .jit import array_to_V_hf, hjit, gjit, vjit from .math.linalg import ( div_Vs_hf, matmul_VV_hf, @@ -24,8 +25,9 @@ __all__ = [ "eccentricity_vector_hf", "eccentricity_vector_gf", + "circular_velocity_hf", + "circular_velocity_vf", # TODO - "circular_velocity", "rv_pqw", "coe_rotation_matrix", "coe2rv", @@ -70,8 +72,8 @@ def eccentricity_vector_gf(k, r, v, e): e[0], e[1], e[2] = eccentricity_vector_hf(k, array_to_V_hf(r), array_to_V_hf(v)) -@jit -def circular_velocity(k, a): +@hjit("f(f,f)") +def circular_velocity_hf(k, a): r"""Compute circular velocity for a given body given the gravitational parameter and the semimajor axis. .. math:: @@ -86,7 +88,16 @@ def circular_velocity(k, a): Semimajor Axis """ - return np.sqrt(k / a) + return sqrt(k / a) + + +@vjit("f(f,f)") +def circular_velocity_vf(k, a): + """ + Vectorized circular_velocity + """ + + return circular_velocity_hf(k, a) @jit diff --git a/src/hapsira/core/thrust/change_a_inc.py b/src/hapsira/core/thrust/change_a_inc.py index 351aefab9..cbd166ef2 100644 --- a/src/hapsira/core/thrust/change_a_inc.py +++ b/src/hapsira/core/thrust/change_a_inc.py @@ -2,7 +2,7 @@ import numpy as np from numpy import cross -from hapsira.core.elements import circular_velocity +from hapsira.core.elements import circular_velocity_hf from ..jit import array_to_V_hf from ..math.linalg import norm_hf @@ -37,8 +37,8 @@ def beta_0(V_0, V_f, inc_0, inc_f): @jit def compute_parameters(k, a_0, a_f, inc_0, inc_f): """Compute parameters of the model.""" - V_0 = circular_velocity(k, a_0) - V_f = circular_velocity(k, a_f) + V_0 = circular_velocity_hf(k, a_0) + V_f = circular_velocity_hf(k, a_f) beta_0_ = beta_0(V_0, V_f, inc_0, inc_f) return V_0, V_f, beta_0_ diff --git a/src/hapsira/core/thrust/change_argp.py b/src/hapsira/core/thrust/change_argp.py index a66f59650..192a0eb7b 100644 --- a/src/hapsira/core/thrust/change_argp.py +++ b/src/hapsira/core/thrust/change_argp.py @@ -2,7 +2,7 @@ import numpy as np from numpy import cross -from hapsira.core.elements import circular_velocity, rv2coe +from hapsira.core.elements import circular_velocity_hf, rv2coe from ..jit import array_to_V_hf from ..math.linalg import norm_hf @@ -20,7 +20,7 @@ def delta_V(V, ecc, argp_0, argp_f, f, A): @jit def extra_quantities(k, a, ecc, argp_0, argp_f, f, A=0.0): """Extra quantities given by the model.""" - V = circular_velocity(k, a) + V = circular_velocity_hf(k, a) delta_V_ = delta_V(V, ecc, argp_0, argp_f, f, A) t_f_ = delta_V_ / f diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index 8b3f4be61..50d8d2484 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -10,7 +10,7 @@ from numpy import cross from hapsira.core.elements import ( - circular_velocity, + circular_velocity_vf, eccentricity_vector_gf, rv2coe, ) @@ -85,7 +85,7 @@ def a_d(t0, u_, k_): accel_v = f * (np.cos(beta_) * thrust_unit + np.sin(beta_) * w_) return accel_v - delta_v = delta_V(circular_velocity(k, a), ecc_0, ecc_f, beta_0) + delta_v = delta_V(circular_velocity_vf(k, a), ecc_0, ecc_f, beta_0) t_f = delta_t(delta_v, f) return a_d, delta_v, t_f diff --git a/src/hapsira/core/thrust/change_ecc_quasioptimal.py b/src/hapsira/core/thrust/change_ecc_quasioptimal.py index c78d2487c..d26895522 100644 --- a/src/hapsira/core/thrust/change_ecc_quasioptimal.py +++ b/src/hapsira/core/thrust/change_ecc_quasioptimal.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from hapsira.core.elements import circular_velocity +from hapsira.core.elements import circular_velocity_hf @jit @@ -13,7 +13,7 @@ def delta_V(V_0, ecc_0, ecc_f): @jit def extra_quantities(k, a, ecc_0, ecc_f, f): """Extra quantities given by the model.""" - V_0 = circular_velocity(k, a) + V_0 = circular_velocity_hf(k, a) delta_V_ = delta_V(V_0, ecc_0, ecc_f) t_f_ = delta_V_ / f diff --git a/src/hapsira/twobody/elements.py b/src/hapsira/twobody/elements.py index 3f6f55711..2420fa722 100644 --- a/src/hapsira/twobody/elements.py +++ b/src/hapsira/twobody/elements.py @@ -2,7 +2,7 @@ import numpy as np from hapsira.core.elements import ( - circular_velocity as circular_velocity_fast, + circular_velocity_vf, coe2rv as coe2rv_fast, coe2rv_many as coe2rv_many_fast, eccentricity_vector_gf, @@ -18,7 +18,7 @@ @u.quantity_input(k=u_km3s2, a=u.km) def circular_velocity(k, a): """Circular velocity for a given body (k) and semimajor axis (a).""" - return circular_velocity_fast(k.to_value(u_km3s2), a.to_value(u.km)) * u_kms + return circular_velocity_vf(k.to_value(u_km3s2), a.to_value(u.km)) * u_kms @u.quantity_input(k=u_km3s2, a=u.km) From e1827e862a817e273dd0ffbce89b951fa71e22e6 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 14:47:37 +0100 Subject: [PATCH 049/346] jit rv_pqw --- src/hapsira/core/elements.py | 26 ++++++++++++++++---------- src/hapsira/core/maneuver.py | 14 +++++++------- 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 1c2879dfb..234701080 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -2,12 +2,12 @@ convert between different elements that define the orbit of a body. """ -from math import sqrt +from math import cos, sqrt, sin import sys from numba import njit as jit, prange import numpy as np -from numpy import cos, cross, sin +from numpy import cross from hapsira.core.angles import E_to_nu_hf, F_to_nu_hf from hapsira.core.util import rotation_matrix @@ -27,8 +27,8 @@ "eccentricity_vector_gf", "circular_velocity_hf", "circular_velocity_vf", + "rv_pqw_hf", # TODO - "rv_pqw", "coe_rotation_matrix", "coe2rv", "coe2rv_many", @@ -100,8 +100,8 @@ def circular_velocity_vf(k, a): return circular_velocity_hf(k, a) -@jit -def rv_pqw(k, p, ecc, nu): +@hjit("Tuple([V,V])(f,f,f,f)") +def rv_pqw_hf(k, p, ecc, nu): r"""Returns r and v vectors in perifocal frame. Parameters @@ -155,10 +155,16 @@ def rv_pqw(k, p, ecc, nu): v = [-5753.30180931 -1328.66813933 0] [m]/[s] """ - pqw = np.array([[cos(nu), sin(nu), 0], [-sin(nu), ecc + cos(nu), 0]]) * np.array( - [[p / (1 + ecc * cos(nu))], [sqrt(k / p)]] + + sin_nu = sin(nu) + cos_nu = cos(nu) + a = p / (1 + ecc * cos_nu) + b = sqrt(k / p) + + return ( + (cos_nu * a, sin_nu * a, 0.0), + (-sin_nu * b, (ecc + cos_nu) * b, 0.0), ) - return pqw @jit @@ -225,10 +231,10 @@ def coe2rv(k, p, ecc, inc, raan, argp, nu): \end{bmatrix} """ - pqw = rv_pqw(k, p, ecc, nu) + pqw = rv_pqw_hf(k, p, ecc, nu) rm = coe_rotation_matrix(inc, raan, argp) - ijk = pqw @ rm.T + ijk = np.array(pqw) @ rm.T return ijk diff --git a/src/hapsira/core/maneuver.py b/src/hapsira/core/maneuver.py index 8ea963488..a1de26435 100644 --- a/src/hapsira/core/maneuver.py +++ b/src/hapsira/core/maneuver.py @@ -4,7 +4,7 @@ import numpy as np from numpy import cross -from hapsira.core.elements import coe_rotation_matrix, rv2coe, rv_pqw +from hapsira.core.elements import coe_rotation_matrix, rv2coe, rv_pqw_hf from .jit import array_to_V_hf from .math.linalg import norm_hf @@ -46,10 +46,10 @@ def hohmann(k, rv, r_f): h_i = norm_hf(array_to_V_hf(cross(*rv))) p_i = h_i**2 / k - r_i, v_i = rv_pqw(k, p_i, ecc, nu) + r_i, v_i = rv_pqw_hf(k, p_i, ecc, nu) - r_i = norm_hf(array_to_V_hf(r_i)) - v_i = norm_hf(array_to_V_hf(v_i)) + r_i = norm_hf(r_i) + v_i = norm_hf(v_i) a_trans = (r_i + r_f) / 2 dv_a = np.sqrt(2 * k / r_i - k / a_trans) - v_i @@ -116,10 +116,10 @@ def bielliptic(k, r_b, r_f, rv): h_i = norm_hf(array_to_V_hf(cross(*rv))) p_i = h_i**2 / k - r_i, v_i = rv_pqw(k, p_i, ecc, nu) + r_i, v_i = rv_pqw_hf(k, p_i, ecc, nu) - r_i = norm_hf(array_to_V_hf(r_i)) - v_i = norm_hf(array_to_V_hf(v_i)) + r_i = norm_hf(r_i) + v_i = norm_hf(v_i) a_trans1 = (r_i + r_b) / 2 a_trans2 = (r_b + r_f) / 2 From 70489c51d08901e9aa55ce0855123f25d92eee9c Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 15:36:58 +0100 Subject: [PATCH 050/346] draft, broken --- src/hapsira/core/elements.py | 17 +++++---- src/hapsira/core/events.py | 4 +- src/hapsira/core/maneuver.py | 6 +-- src/hapsira/core/math/linalg.py | 22 +++++++++++ src/hapsira/core/util.py | 60 +++++++++++++++++++++++------- tests/tests_core/test_core_util.py | 36 +++++++++++++----- 6 files changed, 108 insertions(+), 37 deletions(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 234701080..6bddabc3c 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -10,11 +10,12 @@ from numpy import cross from hapsira.core.angles import E_to_nu_hf, F_to_nu_hf -from hapsira.core.util import rotation_matrix +from hapsira.core.util import rotation_matrix_hf from .jit import array_to_V_hf, hjit, gjit, vjit from .math.linalg import ( div_Vs_hf, + matmul_MM_hf, matmul_VV_hf, mul_Vs_hf, norm_hf, @@ -28,8 +29,8 @@ "circular_velocity_hf", "circular_velocity_vf", "rv_pqw_hf", + "coe_rotation_matrix_hf", # TODO - "coe_rotation_matrix", "coe2rv", "coe2rv_many", "coe2mee", @@ -167,12 +168,12 @@ def rv_pqw_hf(k, p, ecc, nu): ) -@jit -def coe_rotation_matrix(inc, raan, argp): +@hjit("M(f,f,f)") +def coe_rotation_matrix_hf(inc, raan, argp): """Create a rotation matrix for coe transformation.""" - r = rotation_matrix(raan, 2) - r = r @ rotation_matrix(inc, 0) - r = r @ rotation_matrix(argp, 2) + r = rotation_matrix_hf(raan, 2) + r = matmul_MM_hf(r, rotation_matrix_hf(inc, 0)) + r = matmul_MM_hf(r, rotation_matrix_hf(argp, 2)) return r @@ -232,7 +233,7 @@ def coe2rv(k, p, ecc, inc, raan, argp, nu): """ pqw = rv_pqw_hf(k, p, ecc, nu) - rm = coe_rotation_matrix(inc, raan, argp) + rm = np.array(coe_rotation_matrix_hf(inc, raan, argp)) ijk = np.array(pqw) @ rm.T diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index b661de935..dff6a8d3f 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from hapsira.core.elements import coe_rotation_matrix, rv2coe +from hapsira.core.elements import coe_rotation_matrix_hf, rv2coe from hapsira.core.util import planetocentric_to_AltAz from .jit import array_to_V_hf @@ -38,7 +38,7 @@ def eclipse_function(k, u_, r_sec, R_sec, R_primary, umbra=True): pm = 1 if umbra else -1 p, ecc, inc, raan, argp, nu = rv2coe(k, u_[:3], u_[3:]) - PQW = coe_rotation_matrix(inc, raan, argp) + PQW = np.array(coe_rotation_matrix_hf(inc, raan, argp)) # Make arrays contiguous for faster dot product with numba. P_, Q_ = np.ascontiguousarray(PQW[:, 0]), np.ascontiguousarray(PQW[:, 1]) diff --git a/src/hapsira/core/maneuver.py b/src/hapsira/core/maneuver.py index a1de26435..2fb6ce2a4 100644 --- a/src/hapsira/core/maneuver.py +++ b/src/hapsira/core/maneuver.py @@ -4,7 +4,7 @@ import numpy as np from numpy import cross -from hapsira.core.elements import coe_rotation_matrix, rv2coe, rv_pqw_hf +from hapsira.core.elements import coe_rotation_matrix_hf, rv2coe, rv_pqw_hf from .jit import array_to_V_hf from .math.linalg import norm_hf @@ -58,7 +58,7 @@ def hohmann(k, rv, r_f): dv_a = np.array([0, dv_a, 0]) dv_b = np.array([0, -dv_b, 0]) - rot_matrix = coe_rotation_matrix(inc, raan, argp) + rot_matrix = np.array(coe_rotation_matrix_hf(inc, raan, argp)) dv_a = rot_matrix @ dv_a dv_b = rot_matrix @ dv_b @@ -131,7 +131,7 @@ def bielliptic(k, r_b, r_f, rv): dv_b = np.array([0, -dv_b, 0]) dv_c = np.array([0, dv_c, 0]) - rot_matrix = coe_rotation_matrix(inc, raan, argp) + rot_matrix = np.array(coe_rotation_matrix_hf(inc, raan, argp)) dv_a = rot_matrix @ dv_a dv_b = rot_matrix @ dv_b diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 309996d78..6120be9cc 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -4,6 +4,7 @@ __all__ = [ "div_Vs_hf", + "matmul_MM_hf", "matmul_VV_hf", "mul_Vs_hf", "norm_hf", @@ -17,6 +18,27 @@ def div_Vs_hf(v, s): return v[0] / s, v[1] / s, v[2] / s +@hjit("M(M,M)") +def matmul_MM_hf(a, b): + return ( + ( + a[0][0] * b[0][0] + a[0][1] * b[1][0] + a[0][2] * b[2][0], + a[0][0] * b[0][1] + a[0][1] * b[1][1] + a[0][2] * b[2][1], + a[0][0] * b[0][2] + a[0][1] * b[1][2] + a[0][2] * b[2][2], + ), + ( + a[1][0] * b[0][0] + a[1][1] * b[1][0] + a[1][2] * b[2][0], + a[1][0] * b[0][1] + a[1][1] * b[1][1] + a[1][2] * b[2][1], + a[1][0] * b[0][2] + a[1][1] * b[1][2] + a[1][2] * b[2][2], + ), + ( + a[2][0] * b[0][0] + a[2][1] * b[1][0] + a[2][2] * b[2][0], + a[2][0] * b[0][1] + a[2][1] * b[1][1] + a[2][2] * b[2][1], + a[2][0] * b[0][2] + a[2][1] * b[1][2] + a[2][2] * b[2][2], + ), + ) + + @hjit("V(V,f)") def mul_Vs_hf(v, s): return v[0] * s, v[1] * s, v[2] * s diff --git a/src/hapsira/core/util.py b/src/hapsira/core/util.py index 8689cc0f9..98e044b1f 100644 --- a/src/hapsira/core/util.py +++ b/src/hapsira/core/util.py @@ -1,24 +1,56 @@ +from math import cos, sin + from numba import njit as jit import numpy as np -from numpy import cos, sin +from .jit import hjit, gjit -@jit -def rotation_matrix(angle, axis): - assert axis in (0, 1, 2) - angle = np.asarray(angle) + +__all__ = [ + "rotation_matrix_hf", + "rotation_matrix_gf", + "alinspace", + "spherical_to_cartesian", + "planetocentric_to_AltAz", +] + + +@hjit("M(f,u1)") +def rotation_matrix_hf(angle, axis): c = cos(angle) s = sin(angle) + if axis == 0: + return ( + (1.0, 0.0, 0.0), + (0.0, c, -s), + (0.0, s, c), + ) + if axis == 1: + return ( + (c, 0.0, s), + (0.0, 1.0, 0.0), + (s, 0.0, c), + ) + if axis == 2: + return ( + (c, -s, 0.0), + (s, c, 0.0), + (0.0, 0.0, 1.0), + ) + raise ValueError("Invalid axis: must be one of 0, 1 or 2") + + +@gjit("void(f,u1,f[:,:])", "(),()->(3,3)") +def rotation_matrix_gf(angle, axis, r): + """ + Vectorized rotation_matrix + """ - a1 = (axis + 1) % 3 - a2 = (axis + 2) % 3 - R = np.zeros(angle.shape + (3, 3)) - R[..., axis, axis] = 1.0 - R[..., a1, a1] = c - R[..., a1, a2] = -s - R[..., a2, a1] = s - R[..., a2, a2] = c - return R + ( + (r[0, 0], r[0, 1], r[0, 2]), + (r[1, 0], r[1, 1], r[1, 2]), + (r[2, 0], r[2, 1], r[2, 2]), + ) = rotation_matrix_hf(angle, axis) @jit diff --git a/tests/tests_core/test_core_util.py b/tests/tests_core/test_core_util.py index cc613dc58..b17402c73 100644 --- a/tests/tests_core/test_core_util.py +++ b/tests/tests_core/test_core_util.py @@ -10,20 +10,23 @@ from hapsira.core.util import ( alinspace, - rotation_matrix as rotation_matrix_hapsira, + rotation_matrix_gf, spherical_to_cartesian, ) def _test_rotation_matrix_with_v(v, angle, axis): exp = rotation_matrix_astropy(np.degrees(-angle), "xyz"[axis]) @ v - res = rotation_matrix_hapsira(angle, axis) @ v + res = np.zeros(exp.shape, dtype=exp.dtype) + rotation_matrix_gf(angle, axis, res) + res = res @ v assert_allclose(exp, res) def _test_rotation_matrix(angle, axis): expected = rotation_matrix_astropy(-np.rad2deg(angle), "xyz"[axis]) - result = rotation_matrix_hapsira(angle, axis) + result = np.zeros(expected.shape, dtype=expected.dtype) + rotation_matrix_gf(angle, axis, result) assert_allclose(expected, result) @@ -37,23 +40,36 @@ def test_rotation_matrix(): # These tests are adapted from astropy: # https://github.com/astropy/astropy/blob/main/astropy/coordinates/tests/test_matrix_utilities.py def test_rotation_matrix_astropy(): - assert_array_equal(rotation_matrix_hapsira(0, 0), np.eye(3)) + exp = np.eye(3) + res = np.zeros(exp.shape, dtype=exp.dtype) + rotation_matrix_gf(0, 0, res) + assert_array_equal(res, exp) + + exp = np.array([[0, 0, -1], [0, 1, 0], [1, 0, 0]], dtype=float) + res = np.zeros(exp.shape, dtype=exp.dtype) + rotation_matrix_gf(np.deg2rad(-90), 1, res) assert_allclose( - rotation_matrix_hapsira(np.deg2rad(-90), 1), - [[0, 0, -1], [0, 1, 0], [1, 0, 0]], + res, + exp, atol=1e-12, ) + exp = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=float) + res = np.zeros(exp.shape, dtype=exp.dtype) + rotation_matrix_gf(np.deg2rad(90), 2, res) assert_allclose( - rotation_matrix_hapsira(np.deg2rad(90), 2), - [[0, -1, 0], [1, 0, 0], [0, 0, 1]], + res, + exp, atol=1e-12, ) # make sure it also works for very small angles + exp = rotation_matrix_astropy(-0.000001, "x") + res = np.zeros(exp.shape, dtype=exp.dtype) + rotation_matrix_gf(np.deg2rad(0.000001), 0, res) assert_allclose( - rotation_matrix_astropy(-0.000001, "x"), - rotation_matrix_hapsira(np.deg2rad(0.000001), 0), + exp, + res, ) From 4a1ea04f71237b154fbf4e4600d0f04594647efe Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 16:29:37 +0100 Subject: [PATCH 051/346] fix sign error; fix type for ints --- src/hapsira/core/util.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/hapsira/core/util.py b/src/hapsira/core/util.py index 98e044b1f..96a42cb54 100644 --- a/src/hapsira/core/util.py +++ b/src/hapsira/core/util.py @@ -15,7 +15,7 @@ ] -@hjit("M(f,u1)") +@hjit("M(f,i8)") def rotation_matrix_hf(angle, axis): c = cos(angle) s = sin(angle) @@ -29,7 +29,7 @@ def rotation_matrix_hf(angle, axis): return ( (c, 0.0, s), (0.0, 1.0, 0.0), - (s, 0.0, c), + (-s, 0.0, c), ) if axis == 2: return ( @@ -40,12 +40,14 @@ def rotation_matrix_hf(angle, axis): raise ValueError("Invalid axis: must be one of 0, 1 or 2") -@gjit("void(f,u1,f[:,:])", "(),()->(3,3)") -def rotation_matrix_gf(angle, axis, r): +@gjit("void(f,i8,u1[:],f[:,:])", "(),(),(n)->(n,n)") +def rotation_matrix_gf(angle, axis, dummy, r): """ Vectorized rotation_matrix - """ + `dummy` because of https://github.com/numba/numba/issues/2797 + """ + assert dummy.shape == (3,) ( (r[0, 0], r[0, 1], r[0, 2]), (r[1, 0], r[1, 1], r[1, 2]), From 7f0b565af493bf5c1356cbe47c1dffe0fd988ed7 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 16:29:58 +0100 Subject: [PATCH 052/346] partial test fix, still failing --- tests/tests_core/test_core_util.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/tests_core/test_core_util.py b/tests/tests_core/test_core_util.py index b17402c73..858be0f95 100644 --- a/tests/tests_core/test_core_util.py +++ b/tests/tests_core/test_core_util.py @@ -17,16 +17,16 @@ def _test_rotation_matrix_with_v(v, angle, axis): exp = rotation_matrix_astropy(np.degrees(-angle), "xyz"[axis]) @ v - res = np.zeros(exp.shape, dtype=exp.dtype) - rotation_matrix_gf(angle, axis, res) + res = np.zeros((3, 3), dtype=exp.dtype) + rotation_matrix_gf(angle, axis, np.zeros((3,), dtype="u1"), res) res = res @ v assert_allclose(exp, res) def _test_rotation_matrix(angle, axis): expected = rotation_matrix_astropy(-np.rad2deg(angle), "xyz"[axis]) - result = np.zeros(expected.shape, dtype=expected.dtype) - rotation_matrix_gf(angle, axis, result) + result = np.zeros((3, 3), dtype=expected.dtype) + rotation_matrix_gf(angle, axis, np.zeros((3,), dtype="u1"), result) assert_allclose(expected, result) @@ -41,13 +41,13 @@ def test_rotation_matrix(): # https://github.com/astropy/astropy/blob/main/astropy/coordinates/tests/test_matrix_utilities.py def test_rotation_matrix_astropy(): exp = np.eye(3) - res = np.zeros(exp.shape, dtype=exp.dtype) - rotation_matrix_gf(0, 0, res) + res = np.zeros((3, 3), dtype=exp.dtype) + rotation_matrix_gf(0, 0, np.zeros((3,), dtype="u1"), res) assert_array_equal(res, exp) exp = np.array([[0, 0, -1], [0, 1, 0], [1, 0, 0]], dtype=float) - res = np.zeros(exp.shape, dtype=exp.dtype) - rotation_matrix_gf(np.deg2rad(-90), 1, res) + res = np.zeros((3, 3), dtype=exp.dtype) + rotation_matrix_gf(np.deg2rad(-90), 1, np.zeros((3,), dtype="u1"), res) assert_allclose( res, exp, @@ -55,8 +55,8 @@ def test_rotation_matrix_astropy(): ) exp = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=float) - res = np.zeros(exp.shape, dtype=exp.dtype) - rotation_matrix_gf(np.deg2rad(90), 2, res) + res = np.zeros((3, 3), dtype=exp.dtype) + rotation_matrix_gf(np.deg2rad(90), 2, np.zeros((3,), dtype="u1"), res) assert_allclose( res, exp, @@ -65,8 +65,8 @@ def test_rotation_matrix_astropy(): # make sure it also works for very small angles exp = rotation_matrix_astropy(-0.000001, "x") - res = np.zeros(exp.shape, dtype=exp.dtype) - rotation_matrix_gf(np.deg2rad(0.000001), 0, res) + res = np.zeros((3, 3), dtype=exp.dtype) + rotation_matrix_gf(np.deg2rad(0.000001), 0, np.zeros((3,), dtype="u1"), res) assert_allclose( exp, res, From e787010b91fea0164279749ab7b4ef2a660ab696 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 20:52:27 +0100 Subject: [PATCH 053/346] fix test shape --- tests/tests_core/test_core_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/tests_core/test_core_util.py b/tests/tests_core/test_core_util.py index 858be0f95..846b018f4 100644 --- a/tests/tests_core/test_core_util.py +++ b/tests/tests_core/test_core_util.py @@ -17,7 +17,7 @@ def _test_rotation_matrix_with_v(v, angle, axis): exp = rotation_matrix_astropy(np.degrees(-angle), "xyz"[axis]) @ v - res = np.zeros((3, 3), dtype=exp.dtype) + res = np.zeros(np.asarray(angle).shape + (3, 3), dtype=exp.dtype) rotation_matrix_gf(angle, axis, np.zeros((3,), dtype="u1"), res) res = res @ v assert_allclose(exp, res) From 83a3ae62300f8542dadccc3bc3fee79c48839557 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 23:48:37 +0100 Subject: [PATCH 054/346] more jited operators --- src/hapsira/core/math/linalg.py | 26 +++++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 6120be9cc..8fc331bcd 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -5,11 +5,13 @@ __all__ = [ "div_Vs_hf", "matmul_MM_hf", + "matmul_VM_hf", "matmul_VV_hf", "mul_Vs_hf", "norm_hf", "norm_vf", "sub_VV_hf", + "transpose_M_hf", ] @@ -39,9 +41,13 @@ def matmul_MM_hf(a, b): ) -@hjit("V(V,f)") -def mul_Vs_hf(v, s): - return v[0] * s, v[1] * s, v[2] * s +@hjit("V(V,M)") +def matmul_VM_hf(a, b): + return ( + a[0] * b[0][0] + a[1] * b[1][0] + a[2] * b[2][0], + a[0] * b[0][1] + a[1] * b[1][1] + a[2] * b[2][1], + a[0] * b[0][2] + a[1] * b[1][2] + a[2] * b[2][2], + ) @hjit("f(V,V)") @@ -49,6 +55,11 @@ def matmul_VV_hf(a, b): return a[0] * b[0] + a[1] * b[1] + a[2] * b[2] +@hjit("V(V,f)") +def mul_Vs_hf(v, s): + return v[0] * s, v[1] * s, v[2] * s + + @hjit("f(V)") def norm_hf(a): return sqrt(matmul_VV_hf(a, a)) @@ -63,3 +74,12 @@ def norm_vf(a, b, c): @hjit("V(V,V)") def sub_VV_hf(va, vb): return va[0] - vb[0], va[1] - vb[1], va[2] - vb[2] + + +@hjit("M(M)") +def transpose_M_hf(a): + return ( + (a[0][0], a[1][0], a[2][0]), + (a[0][1], a[1][1], a[2][1]), + (a[0][2], a[1][2], a[2][2]), + ) From 1156335c4eb2f17caf03c3dc1048fcaaba345f52 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 1 Jan 2024 23:49:19 +0100 Subject: [PATCH 055/346] jit coe2rv --- src/hapsira/core/elements.py | 46 +++++++++++----------- src/hapsira/core/propagation/danby.py | 4 +- src/hapsira/core/propagation/farnocchia.py | 4 +- src/hapsira/core/propagation/gooding.py | 4 +- src/hapsira/core/propagation/markley.py | 4 +- src/hapsira/core/propagation/mikkola.py | 4 +- src/hapsira/core/propagation/pimienta.py | 4 +- src/hapsira/core/propagation/recseries.py | 4 +- src/hapsira/twobody/elements.py | 37 ++++++++--------- src/hapsira/twobody/sampling.py | 4 +- src/hapsira/twobody/states.py | 15 ++++++- tests/tests_twobody/test_angles.py | 38 +++++++++++++++--- 12 files changed, 99 insertions(+), 69 deletions(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 6bddabc3c..0a6e4e2a8 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -3,9 +3,8 @@ """ from math import cos, sqrt, sin -import sys -from numba import njit as jit, prange +from numba import njit as jit import numpy as np from numpy import cross @@ -16,10 +15,12 @@ from .math.linalg import ( div_Vs_hf, matmul_MM_hf, + matmul_VM_hf, matmul_VV_hf, mul_Vs_hf, norm_hf, sub_VV_hf, + transpose_M_hf, ) @@ -30,9 +31,9 @@ "circular_velocity_vf", "rv_pqw_hf", "coe_rotation_matrix_hf", + "coe2rv_hf", + "coe2rv_gf", # TODO - "coe2rv", - "coe2rv_many", "coe2mee", "rv2coe", "mee2coe", @@ -177,8 +178,8 @@ def coe_rotation_matrix_hf(inc, raan, argp): return r -@jit -def coe2rv(k, p, ecc, inc, raan, argp, nu): +@hjit("Tuple([V,V])(f,f,f,f,f,f,f)") +def coe2rv_hf(k, p, ecc, inc, raan, argp, nu): r"""Converts from classical orbital to state vectors. Classical orbital elements are converted into position and velocity @@ -204,9 +205,9 @@ def coe2rv(k, p, ecc, inc, raan, argp, nu): Returns ------- - r_ijk: numpy.ndarray + r_ijk: tuple[float,float,float] Position vector in basis ijk. - v_ijk: numpy.ndarray + v_ijk: tuple[float,float,float] Velocity vector in basis ijk. Notes @@ -232,26 +233,23 @@ def coe2rv(k, p, ecc, inc, raan, argp, nu): \end{bmatrix} """ - pqw = rv_pqw_hf(k, p, ecc, nu) - rm = np.array(coe_rotation_matrix_hf(inc, raan, argp)) - - ijk = np.array(pqw) @ rm.T - - return ijk + r, v = rv_pqw_hf(k, p, ecc, nu) + rm = transpose_M_hf(coe_rotation_matrix_hf(inc, raan, argp)) + return matmul_VM_hf(r, rm), matmul_VM_hf(v, rm) -@jit(parallel=sys.maxsize > 2**31) -def coe2rv_many(k, p, ecc, inc, raan, argp, nu): - """Parallel version of coe2rv.""" - n = nu.shape[0] - rr = np.zeros((n, 3)) - vv = np.zeros((n, 3)) +@gjit("void(f,f,f,f,f,f,f,u1[:],f[:],f[:])", "(),(),(),(),(),(),(),(n)->(n),(n)") +def coe2rv_gf(k, p, ecc, inc, raan, argp, nu, dummy, rr, vv): + """ + Vectorized coe2rv - # Disabling pylint warning, see https://github.com/PyCQA/pylint/issues/2910 - for i in prange(n): # pylint: disable=not-an-iterable - rr[i, :], vv[i, :] = coe2rv(k[i], p[i], ecc[i], inc[i], raan[i], argp[i], nu[i]) + `dummy` because of https://github.com/numba/numba/issues/2797 + """ + assert dummy.shape == (3,) - return rr, vv + (rr[0], rr[1], rr[2]), (vv[0], vv[1], vv[2]) = coe2rv_hf( + k, p, ecc, inc, raan, argp, nu + ) @jit diff --git a/src/hapsira/core/propagation/danby.py b/src/hapsira/core/propagation/danby.py index 0f685b12b..080b713b5 100644 --- a/src/hapsira/core/propagation/danby.py +++ b/src/hapsira/core/propagation/danby.py @@ -2,7 +2,7 @@ import numpy as np from hapsira.core.angles import E_to_M_hf, F_to_M_hf, nu_to_E_hf, nu_to_F_hf -from hapsira.core.elements import coe2rv, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe @jit @@ -107,4 +107,4 @@ def danby(k, r0, v0, tof, numiter=20, rtol=1e-8): p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0) nu = danby_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol) - return coe2rv(k, p, ecc, inc, raan, argp, nu) + return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/farnocchia.py b/src/hapsira/core/propagation/farnocchia.py index dc19dc2bd..73de00007 100644 --- a/src/hapsira/core/propagation/farnocchia.py +++ b/src/hapsira/core/propagation/farnocchia.py @@ -15,7 +15,7 @@ nu_to_E_hf, nu_to_F_hf, ) -from hapsira.core.elements import coe2rv, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe @jit @@ -332,4 +332,4 @@ def farnocchia_rv(k, r0, v0, tof): p, ecc, inc, raan, argp, nu0 = rv2coe(k, r0, v0) nu = farnocchia_coe(k, p, ecc, inc, raan, argp, nu0, tof) - return coe2rv(k, p, ecc, inc, raan, argp, nu) + return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/gooding.py b/src/hapsira/core/propagation/gooding.py index 61c737a7f..42818209f 100644 --- a/src/hapsira/core/propagation/gooding.py +++ b/src/hapsira/core/propagation/gooding.py @@ -2,7 +2,7 @@ import numpy as np from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf -from hapsira.core.elements import coe2rv, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe @jit @@ -72,4 +72,4 @@ def gooding(k, r0, v0, tof, numiter=150, rtol=1e-8): p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0) nu = gooding_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol) - return coe2rv(k, p, ecc, inc, raan, argp, nu) + return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/markley.py b/src/hapsira/core/propagation/markley.py index 32ec6bd0c..4412d3682 100644 --- a/src/hapsira/core/propagation/markley.py +++ b/src/hapsira/core/propagation/markley.py @@ -8,7 +8,7 @@ kepler_equation_prime_hf, nu_to_E_hf, ) -from hapsira.core.elements import coe2rv, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe @jit @@ -91,4 +91,4 @@ def markley(k, r0, v0, tof): p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0) nu = markley_coe(k, p, ecc, inc, raan, argp, nu, tof) - return coe2rv(k, p, ecc, inc, raan, argp, nu) + return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/mikkola.py b/src/hapsira/core/propagation/mikkola.py index 5daa5acb8..93aee042e 100644 --- a/src/hapsira/core/propagation/mikkola.py +++ b/src/hapsira/core/propagation/mikkola.py @@ -10,7 +10,7 @@ nu_to_E_hf, nu_to_F_hf, ) -from hapsira.core.elements import coe2rv, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe @jit @@ -125,4 +125,4 @@ def mikkola(k, r0, v0, tof, rtol=None): p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0) nu = mikkola_coe(k, p, ecc, inc, raan, argp, nu, tof) - return coe2rv(k, p, ecc, inc, raan, argp, nu) + return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/pimienta.py b/src/hapsira/core/propagation/pimienta.py index 13e2c492b..5295cec4c 100644 --- a/src/hapsira/core/propagation/pimienta.py +++ b/src/hapsira/core/propagation/pimienta.py @@ -2,7 +2,7 @@ import numpy as np from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf -from hapsira.core.elements import coe2rv, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe @jit @@ -371,4 +371,4 @@ def pimienta(k, r0, v0, tof): p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0) nu = pimienta_coe(k, p, ecc, inc, raan, argp, nu, tof) - return coe2rv(k, p, ecc, inc, raan, argp, nu) + return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/recseries.py b/src/hapsira/core/propagation/recseries.py index 55d818dd6..e94eebccd 100644 --- a/src/hapsira/core/propagation/recseries.py +++ b/src/hapsira/core/propagation/recseries.py @@ -2,7 +2,7 @@ import numpy as np from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf -from hapsira.core.elements import coe2rv, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe @jit @@ -117,4 +117,4 @@ def recseries(k, r0, v0, tof, method="rtol", order=8, numiter=100, rtol=1e-8): k, p, ecc, inc, raan, argp, nu, tof, method, order, numiter, rtol ) - return coe2rv(k, p, ecc, inc, raan, argp, nu) + return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/twobody/elements.py b/src/hapsira/twobody/elements.py index 2420fa722..72bb083cb 100644 --- a/src/hapsira/twobody/elements.py +++ b/src/hapsira/twobody/elements.py @@ -3,8 +3,7 @@ from hapsira.core.elements import ( circular_velocity_vf, - coe2rv as coe2rv_fast, - coe2rv_many as coe2rv_many_fast, + coe2rv_gf, eccentricity_vector_gf, ) from hapsira.core.propagation.farnocchia import ( @@ -190,8 +189,18 @@ def get_eccentricity_critical_inc(ecc=None): return ecc -def coe2rv(k, p, ecc, inc, raan, argp, nu): - rr, vv = coe2rv_fast( +def coe2rv(k, p, ecc, inc, raan, argp, nu, rr=None, vv=None): + """ + TODO document function + + Function works on scalars and arrays + """ + + if rr is None and vv is None: + rr = np.zeros(k.shape + (3,), dtype=k.dtype) + vv = np.zeros(k.shape + (3,), dtype=k.dtype) + + coe2rv_gf( k.to_value(u_km3s2), p.to_value(u.km), ecc.to_value(u.one), @@ -199,26 +208,12 @@ def coe2rv(k, p, ecc, inc, raan, argp, nu): raan.to_value(u.rad), argp.to_value(u.rad), nu.to_value(u.rad), + np.zeros((3,), dtype="u1"), # dummy + rr, + vv, ) rr = rr << u.km vv = vv << (u.km / u.s) return rr, vv - - -def coe2rv_many(k_arr, p_arr, ecc_arr, inc_arr, raan_arr, argp_arr, nu_arr): - rr_arr, vv_arr = coe2rv_many_fast( - k_arr.to_value(u_km3s2), - p_arr.to_value(u.km), - ecc_arr.to_value(u.one), - inc_arr.to_value(u.rad), - raan_arr.to_value(u.rad), - argp_arr.to_value(u.rad), - nu_arr.to_value(u.rad), - ) - - rr_arr = rr_arr << u.km - vv_arr = vv_arr << (u.km / u.s) - - return rr_arr, vv_arr diff --git a/src/hapsira/twobody/sampling.py b/src/hapsira/twobody/sampling.py index 89f4bcf73..cd7a412f9 100644 --- a/src/hapsira/twobody/sampling.py +++ b/src/hapsira/twobody/sampling.py @@ -3,7 +3,7 @@ import numpy as np from hapsira.twobody.angles import E_to_nu, nu_to_E -from hapsira.twobody.elements import coe2rv_many, hyp_nu_limit, t_p +from hapsira.twobody.elements import coe2rv, hyp_nu_limit, t_p from hapsira.twobody.propagation import FarnocchiaPropagator from hapsira.util import alinspace, wrap_angle @@ -144,7 +144,7 @@ def sample(self, orbit): epochs = orbit.epoch + (delta_ts - orbit.t_p) n = nu_values.shape[0] - rr, vv = coe2rv_many( + rr, vv = coe2rv( np.tile(orbit.attractor.k, n), np.tile(orbit.p, n), np.tile(orbit.ecc, n), diff --git a/src/hapsira/twobody/states.py b/src/hapsira/twobody/states.py index 53899ca37..1c51291b1 100644 --- a/src/hapsira/twobody/states.py +++ b/src/hapsira/twobody/states.py @@ -1,8 +1,9 @@ from functools import cached_property from astropy import units as u +import numpy as np -from hapsira.core.elements import coe2mee, coe2rv, mee2coe, mee2rv, rv2coe +from hapsira.core.elements import coe2mee, coe2rv_gf, mee2coe, mee2rv, rv2coe from hapsira.twobody.elements import mean_motion, period, t_p @@ -171,7 +172,17 @@ def to_value(self): def to_vectors(self): """Converts to position and velocity vector representation.""" - r, v = coe2rv(self.attractor.k.to_value(u.km**3 / u.s**2), *self.to_value()) + + r = np.zeros(self.attractor.k.shape + (3,), dtype=self.attractor.k.dtype) + v = np.zeros(self.attractor.k.shape + (3,), dtype=self.attractor.k.dtype) + + coe2rv_gf( + self.attractor.k.to_value(u.km**3 / u.s**2), + *self.to_value(), + np.zeros((3,), dtype="u1"), # dummy + r, + v, + ) return RVState(self.attractor, (r << u.km, v << u.km / u.s), self.plane) diff --git a/tests/tests_twobody/test_angles.py b/tests/tests_twobody/test_angles.py index 45da70fc3..76d7e618a 100644 --- a/tests/tests_twobody/test_angles.py +++ b/tests/tests_twobody/test_angles.py @@ -5,7 +5,7 @@ import pytest from hapsira.bodies import Earth -from hapsira.core.elements import coe2mee, coe2rv, mee2coe, rv2coe +from hapsira.core.elements import coe2mee, coe2rv_gf, mee2coe, rv2coe from hapsira.twobody.angles import ( E_to_M, E_to_nu, @@ -211,7 +211,13 @@ def test_eccentric_to_true_range(E, ecc): def test_convert_between_coe_and_rv_is_transitive(classical): k = Earth.k.to(u.km**3 / u.s**2).value # u.km**3 / u.s**2 - res = rv2coe(k, *coe2rv(k, *classical)) + expected_res = classical + + r, v = np.zeros((3,), dtype=float), np.zeros((3,), dtype=float) + coe2rv_gf(k, *expected_res, np.zeros((3,), dtype="u1"), r, v) + + res = rv2coe(k, r, v) + assert_allclose(res, classical) @@ -222,23 +228,43 @@ def test_convert_between_coe_and_mee_is_transitive(classical): def test_convert_coe_and_rv_circular(circular): k, expected_res = circular - res = rv2coe(k, *coe2rv(k, *expected_res)) + + r, v = np.zeros((3,), dtype=float), np.zeros((3,), dtype=float) + coe2rv_gf(k, *expected_res, np.zeros((3,), dtype="u1"), r, v) + + res = rv2coe(k, r, v) + assert_allclose(res, expected_res, atol=1e-8) def test_convert_coe_and_rv_hyperbolic(hyperbolic): k, expected_res = hyperbolic - res = rv2coe(k, *coe2rv(k, *expected_res)) + + r, v = np.zeros((3,), dtype=float), np.zeros((3,), dtype=float) + coe2rv_gf(k, *expected_res, np.zeros((3,), dtype="u1"), r, v) + + res = rv2coe(k, r, v) + assert_allclose(res, expected_res, atol=1e-8) def test_convert_coe_and_rv_equatorial(equatorial): k, expected_res = equatorial - res = rv2coe(k, *coe2rv(k, *expected_res)) + + r, v = np.zeros((3,), dtype=float), np.zeros((3,), dtype=float) + coe2rv_gf(k, *expected_res, np.zeros((3,), dtype="u1"), r, v) + + res = rv2coe(k, r, v) + assert_allclose(res, expected_res, atol=1e-8) def test_convert_coe_and_rv_circular_equatorial(circular_equatorial): k, expected_res = circular_equatorial - res = rv2coe(k, *coe2rv(k, *expected_res)) + + r, v = np.zeros((3,), dtype=float), np.zeros((3,), dtype=float) + coe2rv_gf(k, *expected_res, np.zeros((3,), dtype="u1"), r, v) + + res = rv2coe(k, r, v) + assert_allclose(res, expected_res, atol=1e-8) From afdd0dafad57af327c38d00d0d4597c090325944 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 2 Jan 2024 00:26:30 +0100 Subject: [PATCH 056/346] jit coe2mee --- src/hapsira/core/elements.py | 31 +++++++++++++++++++++--------- src/hapsira/twobody/states.py | 10 +++++++--- tests/tests_twobody/test_angles.py | 4 ++-- 3 files changed, 31 insertions(+), 14 deletions(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 0a6e4e2a8..008dd2fbe 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -2,7 +2,7 @@ convert between different elements that define the orbit of a body. """ -from math import cos, sqrt, sin +from math import cos, pi, sin, sqrt, tan from numba import njit as jit import numpy as np @@ -33,8 +33,9 @@ "coe_rotation_matrix_hf", "coe2rv_hf", "coe2rv_gf", + "coe2mee_hf", + "coe2mee_gf", # TODO - "coe2mee", "rv2coe", "mee2coe", "mee2rv", @@ -252,8 +253,8 @@ def coe2rv_gf(k, p, ecc, inc, raan, argp, nu, dummy, rr, vv): ) -@jit -def coe2mee(p, ecc, inc, raan, argp, nu): +@hjit("Tuple([f,f,f,f,f,f])(f,f,f,f,f,f)") +def coe2mee_hf(p, ecc, inc, raan, argp, nu): r"""Converts from classical orbital elements to modified equinoctial orbital elements. The definition of the modified equinoctial orbital elements is taken from [Walker, 1985]. @@ -310,20 +311,32 @@ def coe2mee(p, ecc, inc, raan, argp, nu): \end{align} """ - if inc == np.pi: + if inc == pi: raise ValueError( "Cannot compute modified equinoctial set for 180 degrees orbit inclination due to `h` and `k` singularity." ) lonper = raan + argp - f = ecc * np.cos(lonper) - g = ecc * np.sin(lonper) - h = np.tan(inc / 2) * np.cos(raan) - k = np.tan(inc / 2) * np.sin(raan) + f = ecc * cos(lonper) + g = ecc * sin(lonper) + h = tan(inc / 2) * cos(raan) + k = tan(inc / 2) * sin(raan) L = lonper + nu return p, f, g, h, k, L +@gjit( + "void(f,f,f,f,f,f,f[:],f[:],f[:],f[:],f[:],f[:])", + "(),(),(),(),(),()->(),(),(),(),(),()", +) +def coe2mee_gf(p, ecc, inc, raan, argp, nu, p_, f, g, h, k, L): + """ + Vectorized coe2mee + """ + + p_[0], f[0], g[0], h[0], k[0], L[0] = coe2mee_hf(p, ecc, inc, raan, argp, nu) + + @jit def rv2coe(k, r, v, tol=1e-8): r"""Converts from vectors to classical orbital elements. diff --git a/src/hapsira/twobody/states.py b/src/hapsira/twobody/states.py index 1c51291b1..da66497ae 100644 --- a/src/hapsira/twobody/states.py +++ b/src/hapsira/twobody/states.py @@ -3,7 +3,7 @@ from astropy import units as u import numpy as np -from hapsira.core.elements import coe2mee, coe2rv_gf, mee2coe, mee2rv, rv2coe +from hapsira.core.elements import coe2mee_gf, coe2rv_gf, mee2coe, mee2rv, rv2coe from hapsira.twobody.elements import mean_motion, period, t_p @@ -192,12 +192,16 @@ def to_classical(self): def to_equinoctial(self): """Converts to modified equinoctial elements representation.""" - p, f, g, h, k, L = coe2mee(*self.to_value()) + + p, ecc, inc, raan, argp, nu = self.to_value() + p_, f, g, h, k, L = coe2mee_gf( + p, ecc, inc, raan, argp, nu + ) # pylint: disable=E1120,E0633 return ModifiedEquinoctialState( self.attractor, ( - p << u.km, + p_ << u.km, f << u.rad, g << u.rad, h << u.rad, diff --git a/tests/tests_twobody/test_angles.py b/tests/tests_twobody/test_angles.py index 76d7e618a..ba71361c7 100644 --- a/tests/tests_twobody/test_angles.py +++ b/tests/tests_twobody/test_angles.py @@ -5,7 +5,7 @@ import pytest from hapsira.bodies import Earth -from hapsira.core.elements import coe2mee, coe2rv_gf, mee2coe, rv2coe +from hapsira.core.elements import coe2mee_gf, coe2rv_gf, mee2coe, rv2coe from hapsira.twobody.angles import ( E_to_M, E_to_nu, @@ -222,7 +222,7 @@ def test_convert_between_coe_and_rv_is_transitive(classical): def test_convert_between_coe_and_mee_is_transitive(classical): - res = mee2coe(*coe2mee(*classical)) + res = mee2coe(*coe2mee_gf(*classical)) # pylint: disable=E1133 assert_allclose(res, classical) From 667b0e1daf90f19e487e62d2ee2abebb56d651d9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 2 Jan 2024 00:45:10 +0100 Subject: [PATCH 057/346] leave malloc go guvectorize --- src/hapsira/core/thrust/change_ecc_inc.py | 3 +- src/hapsira/twobody/elements.py | 43 ++++++++++++++--------- src/hapsira/twobody/states.py | 7 ++-- tests/tests_core/test_core_util.py | 28 ++++++++------- 4 files changed, 45 insertions(+), 36 deletions(-) diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index 50d8d2484..4ae4b76c9 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -58,8 +58,7 @@ def delta_t(delta_v, f): def change_ecc_inc(k, a, ecc_0, ecc_f, inc_0, inc_f, argp, r, v, f): # We fix the inertial direction at the beginning if ecc_0 > 0.001: # Arbitrary tolerance - e_vec = np.zeros_like(r) - eccentricity_vector_gf(k, r, v, e_vec) + e_vec = eccentricity_vector_gf(k, r, v) # pylint: disable=E1120 ref_vec = e_vec / ecc_0 else: ref_vec = r / norm_hf(array_to_V_hf(r)) diff --git a/src/hapsira/twobody/elements.py b/src/hapsira/twobody/elements.py index 72bb083cb..c75b6c8c9 100644 --- a/src/hapsira/twobody/elements.py +++ b/src/hapsira/twobody/elements.py @@ -42,8 +42,9 @@ def energy(k, r, v): @u.quantity_input(k=u_km3s2, r=u.km, v=u_kms) def eccentricity_vector(k, r, v): """Eccentricity vector.""" - e = np.zeros(r.shape, dtype=r.dtype) - eccentricity_vector_gf(k.to_value(u_km3s2), r.to_value(u.km), v.to_value(u_kms), e) + e = eccentricity_vector_gf( + k.to_value(u_km3s2), r.to_value(u.km), v.to_value(u_kms) + ) # pylint: disable=E1120 return e << u.one @@ -197,21 +198,29 @@ def coe2rv(k, p, ecc, inc, raan, argp, nu, rr=None, vv=None): """ if rr is None and vv is None: - rr = np.zeros(k.shape + (3,), dtype=k.dtype) - vv = np.zeros(k.shape + (3,), dtype=k.dtype) - - coe2rv_gf( - k.to_value(u_km3s2), - p.to_value(u.km), - ecc.to_value(u.one), - inc.to_value(u.rad), - raan.to_value(u.rad), - argp.to_value(u.rad), - nu.to_value(u.rad), - np.zeros((3,), dtype="u1"), # dummy - rr, - vv, - ) + rr, vv = coe2rv_gf( # pylint: disable=E1120,E0633 + k.to_value(u_km3s2), + p.to_value(u.km), + ecc.to_value(u.one), + inc.to_value(u.rad), + raan.to_value(u.rad), + argp.to_value(u.rad), + nu.to_value(u.rad), + np.zeros((3,), dtype="u1"), # dummy + ) + else: + coe2rv_gf( + k.to_value(u_km3s2), + p.to_value(u.km), + ecc.to_value(u.one), + inc.to_value(u.rad), + raan.to_value(u.rad), + argp.to_value(u.rad), + nu.to_value(u.rad), + np.zeros((3,), dtype="u1"), # dummy + rr, + vv, + ) rr = rr << u.km vv = vv << (u.km / u.s) diff --git a/src/hapsira/twobody/states.py b/src/hapsira/twobody/states.py index da66497ae..996737529 100644 --- a/src/hapsira/twobody/states.py +++ b/src/hapsira/twobody/states.py @@ -193,15 +193,12 @@ def to_classical(self): def to_equinoctial(self): """Converts to modified equinoctial elements representation.""" - p, ecc, inc, raan, argp, nu = self.to_value() - p_, f, g, h, k, L = coe2mee_gf( - p, ecc, inc, raan, argp, nu - ) # pylint: disable=E1120,E0633 + p, f, g, h, k, L = coe2mee_gf(*self.to_value()) # pylint: disable=E1120,E0633 return ModifiedEquinoctialState( self.attractor, ( - p_ << u.km, + p << u.km, f << u.rad, g << u.rad, h << u.rad, diff --git a/tests/tests_core/test_core_util.py b/tests/tests_core/test_core_util.py index 846b018f4..ec65c3430 100644 --- a/tests/tests_core/test_core_util.py +++ b/tests/tests_core/test_core_util.py @@ -17,16 +17,18 @@ def _test_rotation_matrix_with_v(v, angle, axis): exp = rotation_matrix_astropy(np.degrees(-angle), "xyz"[axis]) @ v - res = np.zeros(np.asarray(angle).shape + (3, 3), dtype=exp.dtype) - rotation_matrix_gf(angle, axis, np.zeros((3,), dtype="u1"), res) + res = rotation_matrix_gf( + angle, axis, np.zeros((3,), dtype="u1") + ) # pylint: disable=E1120 res = res @ v assert_allclose(exp, res) def _test_rotation_matrix(angle, axis): expected = rotation_matrix_astropy(-np.rad2deg(angle), "xyz"[axis]) - result = np.zeros((3, 3), dtype=expected.dtype) - rotation_matrix_gf(angle, axis, np.zeros((3,), dtype="u1"), result) + result = rotation_matrix_gf( + angle, axis, np.zeros((3,), dtype="u1") + ) # pylint: disable=E1120 assert_allclose(expected, result) @@ -41,13 +43,13 @@ def test_rotation_matrix(): # https://github.com/astropy/astropy/blob/main/astropy/coordinates/tests/test_matrix_utilities.py def test_rotation_matrix_astropy(): exp = np.eye(3) - res = np.zeros((3, 3), dtype=exp.dtype) - rotation_matrix_gf(0, 0, np.zeros((3,), dtype="u1"), res) + res = rotation_matrix_gf(0, 0, np.zeros((3,), dtype="u1")) # pylint: disable=E1120 assert_array_equal(res, exp) exp = np.array([[0, 0, -1], [0, 1, 0], [1, 0, 0]], dtype=float) - res = np.zeros((3, 3), dtype=exp.dtype) - rotation_matrix_gf(np.deg2rad(-90), 1, np.zeros((3,), dtype="u1"), res) + res = rotation_matrix_gf( + np.deg2rad(-90), 1, np.zeros((3,), dtype="u1") + ) # pylint: disable=E1120 assert_allclose( res, exp, @@ -55,8 +57,9 @@ def test_rotation_matrix_astropy(): ) exp = np.array([[0, -1, 0], [1, 0, 0], [0, 0, 1]], dtype=float) - res = np.zeros((3, 3), dtype=exp.dtype) - rotation_matrix_gf(np.deg2rad(90), 2, np.zeros((3,), dtype="u1"), res) + res = rotation_matrix_gf( + np.deg2rad(90), 2, np.zeros((3,), dtype="u1") + ) # pylint: disable=E1120 assert_allclose( res, exp, @@ -65,8 +68,9 @@ def test_rotation_matrix_astropy(): # make sure it also works for very small angles exp = rotation_matrix_astropy(-0.000001, "x") - res = np.zeros((3, 3), dtype=exp.dtype) - rotation_matrix_gf(np.deg2rad(0.000001), 0, np.zeros((3,), dtype="u1"), res) + res = rotation_matrix_gf( + np.deg2rad(0.000001), 0, np.zeros((3,), dtype="u1") + ) # pylint: disable=E1120 assert_allclose( exp, res, From e8d1ab77bcfd18cff10a87f04f111d6d86c9472a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 2 Jan 2024 14:00:40 +0100 Subject: [PATCH 058/346] jit rv2coe --- ...tural-and-artificial-perturbations.myst.md | 6 +- src/hapsira/core/elements.py | 88 ++++++++++++------- src/hapsira/core/events.py | 6 +- src/hapsira/core/maneuver.py | 15 +++- src/hapsira/core/math/linalg.py | 10 +++ src/hapsira/core/propagation/danby.py | 7 +- src/hapsira/core/propagation/farnocchia.py | 7 +- src/hapsira/core/propagation/gooding.py | 7 +- src/hapsira/core/propagation/markley.py | 7 +- src/hapsira/core/propagation/mikkola.py | 7 +- src/hapsira/core/propagation/pimienta.py | 7 +- src/hapsira/core/propagation/recseries.py | 7 +- src/hapsira/core/thrust/change_argp.py | 4 +- src/hapsira/core/thrust/change_ecc_inc.py | 7 +- src/hapsira/twobody/states.py | 12 ++- tests/tests_twobody/test_angles.py | 12 +-- tests/tests_twobody/test_perturbations.py | 43 ++++++--- tests/tests_twobody/test_propagation.py | 5 +- 18 files changed, 178 insertions(+), 79 deletions(-) diff --git a/docs/source/examples/natural-and-artificial-perturbations.myst.md b/docs/source/examples/natural-and-artificial-perturbations.myst.md index 0f80764f2..4f595b70b 100644 --- a/docs/source/examples/natural-and-artificial-perturbations.myst.md +++ b/docs/source/examples/natural-and-artificial-perturbations.myst.md @@ -24,7 +24,7 @@ from astropy import units as u from hapsira.bodies import Earth, Moon from hapsira.constants import rho0_earth, H0_earth -from hapsira.core.elements import rv2coe +from hapsira.core.elements import rv2coe_gf, RV2COE_TOL from hapsira.core.perturbations import ( atmospheric_drag_exponential, third_body, @@ -161,9 +161,9 @@ rr, vv = orbit.to_ephem( ).rv() # This will be easier to compute when this is solved: -# https://github.com/hapsira/hapsira/issues/380 +# https://github.com/poliastro/poliastro/issues/380 raans = [ - rv2coe(k, r, v)[3] + rv2coe_gf(k, r, v, RV2COE_TOL)[3] for r, v in zip(rr.to_value(u.km), vv.to_value(u.km / u.s)) ] ``` diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 008dd2fbe..9de13a01a 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -2,17 +2,17 @@ convert between different elements that define the orbit of a body. """ -from math import cos, pi, sin, sqrt, tan +from math import acos, atan2, cos, log, pi, sin, sqrt, tan from numba import njit as jit import numpy as np -from numpy import cross from hapsira.core.angles import E_to_nu_hf, F_to_nu_hf from hapsira.core.util import rotation_matrix_hf from .jit import array_to_V_hf, hjit, gjit, vjit from .math.linalg import ( + cross_VV_hf, div_Vs_hf, matmul_MM_hf, matmul_VM_hf, @@ -35,13 +35,18 @@ "coe2rv_gf", "coe2mee_hf", "coe2mee_gf", + "RV2COE_TOL", + "rv2coe_hf", + "rv2coe_gf", # TODO - "rv2coe", "mee2coe", "mee2rv", ] +RV2COE_TOL = 1e-8 + + @hjit("V(f,V,V)") def eccentricity_vector_hf(k, r, v): r"""Eccentricity vector. @@ -337,17 +342,17 @@ def coe2mee_gf(p, ecc, inc, raan, argp, nu, p_, f, g, h, k, L): p_[0], f[0], g[0], h[0], k[0], L[0] = coe2mee_hf(p, ecc, inc, raan, argp, nu) -@jit -def rv2coe(k, r, v, tol=1e-8): +@hjit("Tuple([f,f,f,f,f,f])(f,V,V,f)") +def rv2coe_hf(k, r, v, tol): r"""Converts from vectors to classical orbital elements. Parameters ---------- k : float Standard gravitational parameter (km^3 / s^2) - r : numpy.ndarray + r : tuple[float,float,float] Position vector (km) - v : numpy.ndarray + v : tuple[float,float,float] Velocity vector (km / s) tol : float, optional Tolerance for eccentricity and inclination checks, default to 1e-8 @@ -393,7 +398,7 @@ def rv2coe(k, r, v, tol=1e-8): N &= \sqrt{\vec{N}\cdot\vec{N}} \end{align} - 4. The rigth ascension node is computed: + 4. The right ascension node is computed: .. math:: \Omega = \left\{ \begin{array}{lcc} @@ -442,51 +447,74 @@ def rv2coe(k, r, v, tol=1e-8): nu: 28.445804984192122 [deg] """ - h = cross(r, v) - n = cross([0, 0, 1], h) - e = ((v @ v - k / norm_hf(array_to_V_hf(r))) * r - (r @ v) * v) / k - ecc = norm_hf(array_to_V_hf(e)) - p = (h @ h) / k - inc = np.arccos(h[2] / norm_hf(array_to_V_hf(h))) + h = cross_VV_hf(r, v) + n = cross_VV_hf((0, 0, 1), h) + e = mul_Vs_hf( + sub_VV_hf( + mul_Vs_hf(r, (matmul_VV_hf(v, v) - k / norm_hf(r))), + mul_Vs_hf(v, matmul_VV_hf(r, v)), + ), + 1 / k, + ) + ecc = norm_hf(e) + p = matmul_VV_hf(h, h) / k + inc = acos(h[2] / norm_hf(h)) circular = ecc < tol equatorial = abs(inc) < tol if equatorial and not circular: raan = 0 - argp = np.arctan2(e[1], e[0]) % (2 * np.pi) # Longitude of periapsis - nu = np.arctan2((h @ cross(e, r)) / norm_hf(array_to_V_hf(h)), r @ e) + argp = atan2(e[1], e[0]) % (2 * pi) # Longitude of periapsis + nu = atan2(matmul_VV_hf(h, cross_VV_hf(e, r)) / norm_hf(h), matmul_VV_hf(r, e)) elif not equatorial and circular: - raan = np.arctan2(n[1], n[0]) % (2 * np.pi) + raan = atan2(n[1], n[0]) % (2 * pi) argp = 0 # Argument of latitude - nu = np.arctan2((r @ cross(h, n)) / norm_hf(array_to_V_hf(h)), r @ n) + nu = atan2(matmul_VV_hf(r, cross_VV_hf(h, n)) / norm_hf(h), matmul_VV_hf(r, n)) elif equatorial and circular: raan = 0 argp = 0 - nu = np.arctan2(r[1], r[0]) % (2 * np.pi) # True longitude + nu = atan2(r[1], r[0]) % (2 * pi) # True longitude else: a = p / (1 - (ecc**2)) ka = k * a if a > 0: - e_se = (r @ v) / sqrt(ka) - e_ce = norm_hf(array_to_V_hf(r)) * (v @ v) / k - 1 - nu = E_to_nu_hf(np.arctan2(e_se, e_ce), ecc) + e_se = matmul_VV_hf(r, v) / sqrt(ka) + e_ce = norm_hf(r) * matmul_VV_hf(v, v) / k - 1 + nu = E_to_nu_hf(atan2(e_se, e_ce), ecc) else: - e_sh = (r @ v) / sqrt(-ka) - e_ch = norm_hf(array_to_V_hf(r)) * (norm_hf(array_to_V_hf(v)) ** 2) / k - 1 - nu = F_to_nu_hf(np.log((e_ch + e_sh) / (e_ch - e_sh)) / 2, ecc) + e_sh = matmul_VV_hf(r, v) / sqrt(-ka) + e_ch = norm_hf(r) * (norm_hf(v) ** 2) / k - 1 + nu = F_to_nu_hf(log((e_ch + e_sh) / (e_ch - e_sh)) / 2, ecc) - raan = np.arctan2(n[1], n[0]) % (2 * np.pi) - px = r @ n - py = (r @ cross(h, n)) / norm_hf(array_to_V_hf(h)) - argp = (np.arctan2(py, px) - nu) % (2 * np.pi) + raan = atan2(n[1], n[0]) % (2 * pi) + px = matmul_VV_hf(r, n) + py = matmul_VV_hf(r, cross_VV_hf(h, n)) / norm_hf(h) + argp = (atan2(py, px) - nu) % (2 * pi) - nu = (nu + np.pi) % (2 * np.pi) - np.pi + nu = (nu + pi) % (2 * pi) - pi return p, ecc, inc, raan, argp, nu +@gjit( + "void(f,f[:],f[:],f,f[:],f[:],f[:],f[:],f[:],f[:])", + "(),(n),(n),()->(),(),(),(),(),()", +) +def rv2coe_gf(k, r, v, tol, p, ecc, inc, raan, argp, nu): + """ + Vectorized rv2coe + """ + + p[0], ecc[0], inc[0], raan[0], argp[0], nu[0] = rv2coe_hf( + k, + array_to_V_hf(r), + array_to_V_hf(v), + tol, + ) + + @jit def mee2coe(p, f, g, h, k, L): r"""Converts from modified equinoctial orbital elements to classical diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index dff6a8d3f..f86d3c73c 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from hapsira.core.elements import coe_rotation_matrix_hf, rv2coe +from hapsira.core.elements import coe_rotation_matrix_hf, rv2coe_hf, RV2COE_TOL from hapsira.core.util import planetocentric_to_AltAz from .jit import array_to_V_hf @@ -36,7 +36,9 @@ def eclipse_function(k, u_, r_sec, R_sec, R_primary, umbra=True): """ # Plus or minus condition pm = 1 if umbra else -1 - p, ecc, inc, raan, argp, nu = rv2coe(k, u_[:3], u_[3:]) + p, ecc, inc, raan, argp, nu = rv2coe_hf( + k, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), RV2COE_TOL + ) PQW = np.array(coe_rotation_matrix_hf(inc, raan, argp)) # Make arrays contiguous for faster dot product with numba. diff --git a/src/hapsira/core/maneuver.py b/src/hapsira/core/maneuver.py index 2fb6ce2a4..2b6800331 100644 --- a/src/hapsira/core/maneuver.py +++ b/src/hapsira/core/maneuver.py @@ -4,7 +4,12 @@ import numpy as np from numpy import cross -from hapsira.core.elements import coe_rotation_matrix_hf, rv2coe, rv_pqw_hf +from hapsira.core.elements import ( + coe_rotation_matrix_hf, + rv2coe_hf, + RV2COE_TOL, + rv_pqw_hf, +) from .jit import array_to_V_hf from .math.linalg import norm_hf @@ -42,7 +47,9 @@ def hohmann(k, rv, r_f): Final orbital radius """ - _, ecc, inc, raan, argp, nu = rv2coe(k, *rv) + _, ecc, inc, raan, argp, nu = rv2coe_hf( + k, array_to_V_hf(rv[0]), array_to_V_hf(rv[1]), RV2COE_TOL + ) h_i = norm_hf(array_to_V_hf(cross(*rv))) p_i = h_i**2 / k @@ -112,7 +119,9 @@ def bielliptic(k, r_b, r_f, rv): Position and velocity vectors """ - _, ecc, inc, raan, argp, nu = rv2coe(k, *rv) + _, ecc, inc, raan, argp, nu = rv2coe_hf( + k, array_to_V_hf(rv[0]), array_to_V_hf(rv[1]), RV2COE_TOL + ) h_i = norm_hf(array_to_V_hf(cross(*rv))) p_i = h_i**2 / k diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 8fc331bcd..30f0ce1fc 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -3,6 +3,7 @@ from ..jit import hjit, vjit __all__ = [ + "cross_VV_hf", "div_Vs_hf", "matmul_MM_hf", "matmul_VM_hf", @@ -15,6 +16,15 @@ ] +@hjit("V(V,V)") +def cross_VV_hf(a, b): + return ( + a[1] * b[2] - a[2] * b[1], + a[2] * b[0] - a[0] * b[2], + a[0] * b[1] - a[1] * b[0], + ) + + @hjit("V(V,f)") def div_Vs_hf(v, s): return v[0] / s, v[1] / s, v[2] / s diff --git a/src/hapsira/core/propagation/danby.py b/src/hapsira/core/propagation/danby.py index 080b713b5..b4530797c 100644 --- a/src/hapsira/core/propagation/danby.py +++ b/src/hapsira/core/propagation/danby.py @@ -2,7 +2,8 @@ import numpy as np from hapsira.core.angles import E_to_M_hf, F_to_M_hf, nu_to_E_hf, nu_to_F_hf -from hapsira.core.elements import coe2rv_hf, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf @jit @@ -104,7 +105,9 @@ def danby(k, r0, v0, tof, numiter=20, rtol=1e-8): Equation* with DOI: https://doi.org/10.1007/BF01686811 """ # Solve first for eccentricity and mean anomaly - p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0) + p, ecc, inc, raan, argp, nu = rv2coe_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL + ) nu = danby_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol) return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/farnocchia.py b/src/hapsira/core/propagation/farnocchia.py index 73de00007..9faf5268b 100644 --- a/src/hapsira/core/propagation/farnocchia.py +++ b/src/hapsira/core/propagation/farnocchia.py @@ -15,7 +15,8 @@ nu_to_E_hf, nu_to_F_hf, ) -from hapsira.core.elements import coe2rv_hf, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf @jit @@ -329,7 +330,9 @@ def farnocchia_rv(k, r0, v0, tof): """ # get the initial true anomaly and orbit parameters that are constant over time - p, ecc, inc, raan, argp, nu0 = rv2coe(k, r0, v0) + p, ecc, inc, raan, argp, nu0 = rv2coe_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL + ) nu = farnocchia_coe(k, p, ecc, inc, raan, argp, nu0, tof) return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/gooding.py b/src/hapsira/core/propagation/gooding.py index 42818209f..51c458f5f 100644 --- a/src/hapsira/core/propagation/gooding.py +++ b/src/hapsira/core/propagation/gooding.py @@ -2,7 +2,8 @@ import numpy as np from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf -from hapsira.core.elements import coe2rv_hf, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf @jit @@ -69,7 +70,9 @@ def gooding(k, r0, v0, tof, numiter=150, rtol=1e-8): Original paper for the algorithm: https://doi.org/10.1007/BF01238923 """ # Solve first for eccentricity and mean anomaly - p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0) + p, ecc, inc, raan, argp, nu = rv2coe_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL + ) nu = gooding_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol) return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/markley.py b/src/hapsira/core/propagation/markley.py index 4412d3682..5320b6a01 100644 --- a/src/hapsira/core/propagation/markley.py +++ b/src/hapsira/core/propagation/markley.py @@ -8,7 +8,8 @@ kepler_equation_prime_hf, nu_to_E_hf, ) -from hapsira.core.elements import coe2rv_hf, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf @jit @@ -88,7 +89,9 @@ def markley(k, r0, v0, tof): """ # Solve first for eccentricity and mean anomaly - p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0) + p, ecc, inc, raan, argp, nu = rv2coe_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL + ) nu = markley_coe(k, p, ecc, inc, raan, argp, nu, tof) return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/mikkola.py b/src/hapsira/core/propagation/mikkola.py index 93aee042e..538e9d8d8 100644 --- a/src/hapsira/core/propagation/mikkola.py +++ b/src/hapsira/core/propagation/mikkola.py @@ -10,7 +10,8 @@ nu_to_E_hf, nu_to_F_hf, ) -from hapsira.core.elements import coe2rv_hf, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf @jit @@ -122,7 +123,9 @@ def mikkola(k, r0, v0, tof, rtol=None): Original paper: https://doi.org/10.1007/BF01235850 """ # Solving for the classical elements - p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0) + p, ecc, inc, raan, argp, nu = rv2coe_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL + ) nu = mikkola_coe(k, p, ecc, inc, raan, argp, nu, tof) return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/pimienta.py b/src/hapsira/core/propagation/pimienta.py index 5295cec4c..f468fc07b 100644 --- a/src/hapsira/core/propagation/pimienta.py +++ b/src/hapsira/core/propagation/pimienta.py @@ -2,7 +2,8 @@ import numpy as np from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf -from hapsira.core.elements import coe2rv_hf, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf @jit @@ -368,7 +369,9 @@ def pimienta(k, r0, v0, tof): # TODO: implement hyperbolic case # Solve first for eccentricity and mean anomaly - p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0) + p, ecc, inc, raan, argp, nu = rv2coe_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL + ) nu = pimienta_coe(k, p, ecc, inc, raan, argp, nu, tof) return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) diff --git a/src/hapsira/core/propagation/recseries.py b/src/hapsira/core/propagation/recseries.py index e94eebccd..99ff501d3 100644 --- a/src/hapsira/core/propagation/recseries.py +++ b/src/hapsira/core/propagation/recseries.py @@ -2,7 +2,8 @@ import numpy as np from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf -from hapsira.core.elements import coe2rv_hf, rv2coe +from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf @jit @@ -112,7 +113,9 @@ def recseries(k, r0, v0, tof, method="rtol", order=8, numiter=100, rtol=1e-8): with DOI: http://dx.doi.org/10.13140/RG.2.2.18578.58563/1 """ # Solve first for eccentricity and mean anomaly - p, ecc, inc, raan, argp, nu = rv2coe(k, r0, v0) + p, ecc, inc, raan, argp, nu = rv2coe_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL + ) nu = recseries_coe( k, p, ecc, inc, raan, argp, nu, tof, method, order, numiter, rtol ) diff --git a/src/hapsira/core/thrust/change_argp.py b/src/hapsira/core/thrust/change_argp.py index 192a0eb7b..dbd9a1429 100644 --- a/src/hapsira/core/thrust/change_argp.py +++ b/src/hapsira/core/thrust/change_argp.py @@ -2,7 +2,7 @@ import numpy as np from numpy import cross -from hapsira.core.elements import circular_velocity_hf, rv2coe +from hapsira.core.elements import circular_velocity_hf, rv2coe_hf, RV2COE_TOL from ..jit import array_to_V_hf from ..math.linalg import norm_hf @@ -58,7 +58,7 @@ def change_argp(k, a, ecc, argp_0, argp_f, f): def a_d(t0, u_, k): r = u_[:3] v = u_[3:] - nu = rv2coe(k, r, v)[-1] + nu = rv2coe_hf(k, array_to_V_hf(r), array_to_V_hf(v), RV2COE_TOL)[-1] alpha_ = nu - np.pi / 2 diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index 4ae4b76c9..5218804d6 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -3,7 +3,7 @@ References ---------- * Pollard, J. E. "Simplified Analysis of Low-Thrust Orbital Maneuvers", 2000. - +rv2coe """ from numba import njit as jit import numpy as np @@ -12,7 +12,8 @@ from hapsira.core.elements import ( circular_velocity_vf, eccentricity_vector_gf, - rv2coe, + rv2coe_hf, + RV2COE_TOL, ) from ..jit import array_to_V_hf @@ -73,7 +74,7 @@ def change_ecc_inc(k, a, ecc_0, ecc_f, inc_0, inc_f, argp, r, v, f): def a_d(t0, u_, k_): r_ = u_[:3] v_ = u_[3:] - nu = rv2coe(k_, r_, v_)[-1] + nu = rv2coe_hf(k_, array_to_V_hf(r_), array_to_V_hf(v_), RV2COE_TOL)[-1] beta_ = beta_0 * np.sign( np.cos(nu) ) # The sign of ß reverses at minor axis crossings diff --git a/src/hapsira/twobody/states.py b/src/hapsira/twobody/states.py index 996737529..3b905f418 100644 --- a/src/hapsira/twobody/states.py +++ b/src/hapsira/twobody/states.py @@ -3,7 +3,14 @@ from astropy import units as u import numpy as np -from hapsira.core.elements import coe2mee_gf, coe2rv_gf, mee2coe, mee2rv, rv2coe +from hapsira.core.elements import ( + coe2mee_gf, + coe2rv_gf, + mee2coe, + mee2rv, + rv2coe_gf, + RV2COE_TOL, +) from hapsira.twobody.elements import mean_motion, period, t_p @@ -243,9 +250,10 @@ def to_vectors(self): def to_classical(self): """Converts to classical orbital elements representation.""" - (p, ecc, inc, raan, argp, nu) = rv2coe( + (p, ecc, inc, raan, argp, nu) = rv2coe_gf( # pylint: disable=E1120,E0633 self.attractor.k.to_value(u.km**3 / u.s**2), *self.to_value(), + RV2COE_TOL, ) return ClassicalState( diff --git a/tests/tests_twobody/test_angles.py b/tests/tests_twobody/test_angles.py index ba71361c7..453959c23 100644 --- a/tests/tests_twobody/test_angles.py +++ b/tests/tests_twobody/test_angles.py @@ -5,7 +5,7 @@ import pytest from hapsira.bodies import Earth -from hapsira.core.elements import coe2mee_gf, coe2rv_gf, mee2coe, rv2coe +from hapsira.core.elements import coe2mee_gf, coe2rv_gf, mee2coe, rv2coe_gf, RV2COE_TOL from hapsira.twobody.angles import ( E_to_M, E_to_nu, @@ -216,7 +216,7 @@ def test_convert_between_coe_and_rv_is_transitive(classical): r, v = np.zeros((3,), dtype=float), np.zeros((3,), dtype=float) coe2rv_gf(k, *expected_res, np.zeros((3,), dtype="u1"), r, v) - res = rv2coe(k, r, v) + res = rv2coe_gf(k, r, v, RV2COE_TOL) # pylint: disable=E1120 assert_allclose(res, classical) @@ -232,7 +232,7 @@ def test_convert_coe_and_rv_circular(circular): r, v = np.zeros((3,), dtype=float), np.zeros((3,), dtype=float) coe2rv_gf(k, *expected_res, np.zeros((3,), dtype="u1"), r, v) - res = rv2coe(k, r, v) + res = rv2coe_gf(k, r, v, RV2COE_TOL) # pylint: disable=E1120 assert_allclose(res, expected_res, atol=1e-8) @@ -243,7 +243,7 @@ def test_convert_coe_and_rv_hyperbolic(hyperbolic): r, v = np.zeros((3,), dtype=float), np.zeros((3,), dtype=float) coe2rv_gf(k, *expected_res, np.zeros((3,), dtype="u1"), r, v) - res = rv2coe(k, r, v) + res = rv2coe_gf(k, r, v, RV2COE_TOL) # pylint: disable=E1120 assert_allclose(res, expected_res, atol=1e-8) @@ -254,7 +254,7 @@ def test_convert_coe_and_rv_equatorial(equatorial): r, v = np.zeros((3,), dtype=float), np.zeros((3,), dtype=float) coe2rv_gf(k, *expected_res, np.zeros((3,), dtype="u1"), r, v) - res = rv2coe(k, r, v) + res = rv2coe_gf(k, r, v, RV2COE_TOL) # pylint: disable=E1120 assert_allclose(res, expected_res, atol=1e-8) @@ -265,6 +265,6 @@ def test_convert_coe_and_rv_circular_equatorial(circular_equatorial): r, v = np.zeros((3,), dtype=float), np.zeros((3,), dtype=float) coe2rv_gf(k, *expected_res, np.zeros((3,), dtype="u1"), r, v) - res = rv2coe(k, r, v) + res = rv2coe_gf(k, r, v, RV2COE_TOL) # pylint: disable=E1120 assert_allclose(res, expected_res, atol=1e-8) diff --git a/tests/tests_twobody/test_perturbations.py b/tests/tests_twobody/test_perturbations.py index 630a41b7e..8784fd236 100644 --- a/tests/tests_twobody/test_perturbations.py +++ b/tests/tests_twobody/test_perturbations.py @@ -10,14 +10,14 @@ from hapsira.bodies import Earth, Moon, Sun from hapsira.constants import H0_earth, Wdivc_sun, rho0_earth -from hapsira.core.elements import rv2coe +from hapsira.core.elements import rv2coe_gf, RV2COE_TOL from hapsira.core.perturbations import ( J2_perturbation, J3_perturbation, atmospheric_drag, atmospheric_drag_exponential, radiation_pressure, - third_body, + third_body, # pylint: disable=E1120,E1136 ) from hapsira.core.propagation import func_twobody from hapsira.earth.atmosphere import COESA76 @@ -51,8 +51,12 @@ def f(t0, u_, k): k = Earth.k.to(u.km**3 / u.s**2).value - _, _, _, raan0, argp0, _ = rv2coe(k, r0, v0) - _, _, _, raan, argp, _ = rv2coe(k, rr[0].to(u.km).value, vv[0].to(u.km / u.s).value) + _, _, _, raan0, argp0, _ = rv2coe_gf( # pylint: disable=E1120,E0633 + k, r0, v0, RV2COE_TOL + ) + _, _, _, raan, argp, _ = rv2coe_gf( # pylint: disable=E1120,E0633 + k, rr[0].to(u.km).value, vv[0].to(u.km / u.s).value, RV2COE_TOL + ) raan_variation_rate = (raan - raan0) / tofs[0].to(u.s).value # type: ignore argp_variation_rate = (argp - argp0) / tofs[0].to(u.s).value # type: ignore @@ -148,13 +152,23 @@ def f_combined(t0, u_, k): a_values_J2 = np.array( [ - rv2coe(k, ri, vi)[0] / (1.0 - rv2coe(k, ri, vi)[1] ** 2) + rv2coe_gf(k, ri, vi, RV2COE_TOL)[0] # pylint: disable=E1120,E1136 + / ( + 1.0 + - rv2coe_gf(k, ri, vi, RV2COE_TOL)[1] # pylint: disable=E1120,E1136 + ** 2 + ) for ri, vi in zip(r_J2.to(u.km).value, v_J2.to(u.km / u.s).value) ] ) a_values_J3 = np.array( [ - rv2coe(k, ri, vi)[0] / (1.0 - rv2coe(k, ri, vi)[1] ** 2) + rv2coe_gf(k, ri, vi, RV2COE_TOL)[0] # pylint: disable=E1120,E1136 + / ( + 1.0 + - rv2coe_gf(k, ri, vi, RV2COE_TOL)[1] # pylint: disable=E1120,E1136 + ** 2 + ) for ri, vi in zip(r_J3.to(u.km).value, v_J3.to(u.km / u.s).value) ] ) @@ -162,13 +176,13 @@ def f_combined(t0, u_, k): ecc_values_J2 = np.array( [ - rv2coe(k, ri, vi)[1] + rv2coe_gf(k, ri, vi, RV2COE_TOL)[1] # pylint: disable=E1120,E1136 for ri, vi in zip(r_J2.to(u.km).value, v_J2.to(u.km / u.s).value) ] ) ecc_values_J3 = np.array( [ - rv2coe(k, ri, vi)[1] + rv2coe_gf(k, ri, vi, RV2COE_TOL)[1] # pylint: disable=E1120,E1136 for ri, vi in zip(r_J3.to(u.km).value, v_J3.to(u.km / u.s).value) ] ) @@ -176,13 +190,13 @@ def f_combined(t0, u_, k): inc_values_J2 = np.array( [ - rv2coe(k, ri, vi)[2] + rv2coe_gf(k, ri, vi, RV2COE_TOL)[2] # pylint: disable=E1120,E1136 for ri, vi in zip(r_J2.to(u.km).value, v_J2.to(u.km / u.s).value) ] ) inc_values_J3 = np.array( [ - rv2coe(k, ri, vi)[2] + rv2coe_gf(k, ri, vi, RV2COE_TOL)[2] # pylint: disable=E1120,E1136 for ri, vi in zip(r_J3.to(u.km).value, v_J3.to(u.km / u.s).value) ] ) @@ -573,7 +587,10 @@ def f(t0, u_, k): incs, raans, argps = [], [], [] for ri, vi in zip(rr.to_value(u.km), vv.to_value(u.km / u.s)): angles = Angle( - rv2coe(Earth.k.to_value(u.km**3 / u.s**2), ri, vi)[2:5] * u.rad + rv2coe_gf( # pylint: disable=E1120,E1136 + Earth.k.to_value(u.km**3 / u.s**2), ri, vi, RV2COE_TOL + )[2:5] + * u.rad ) # inc, raan, argp angles = angles.wrap_at(180 * u.deg) incs.append(angles[0].value) @@ -670,7 +687,9 @@ def f(t0, u_, k): delta_eccs, delta_incs, delta_raans, delta_argps = [], [], [], [] for ri, vi in zip(rr.to(u.km).value, vv.to(u.km / u.s).value): - orbit_params = rv2coe(Earth.k.to(u.km**3 / u.s**2).value, ri, vi) + orbit_params = rv2coe_gf( # pylint: disable=E1120,E1136 + Earth.k.to(u.km**3 / u.s**2).value, ri, vi, RV2COE_TOL + ) delta_eccs.append(orbit_params[1] - initial.ecc.value) delta_incs.append((orbit_params[2] * u.rad).to(u.deg).value - initial.inc.value) delta_raans.append( diff --git a/tests/tests_twobody/test_propagation.py b/tests/tests_twobody/test_propagation.py index 0d20983bf..1a8a4bdad 100644 --- a/tests/tests_twobody/test_propagation.py +++ b/tests/tests_twobody/test_propagation.py @@ -9,7 +9,7 @@ from hapsira.bodies import Earth, Moon, Sun from hapsira.constants import J2000 -from hapsira.core.elements import rv2coe +from hapsira.core.elements import rv2coe_gf, RV2COE_TOL from hapsira.core.propagation import func_twobody from hapsira.examples import iss from hapsira.frames import Planes @@ -203,10 +203,11 @@ def test_propagation_parabolic(propagator): orbit = Orbit.parabolic(Earth, p, _a, _a, _a, _a) orbit = orbit.propagate(0.8897 / 2.0 * u.h, method=propagator()) - _, _, _, _, _, nu0 = rv2coe( + _, _, _, _, _, nu0 = rv2coe_gf( # pylint: disable=E1120,E0633 Earth.k.to(u.km**3 / u.s**2).value, orbit.r.to(u.km).value, orbit.v.to(u.km / u.s).value, + RV2COE_TOL, ) assert_quantity_allclose(nu0, np.deg2rad(90.0), rtol=1e-4) From e34058355611a501fde35baf398787230356e6d0 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 2 Jan 2024 18:31:00 +0100 Subject: [PATCH 059/346] jit mee2coe --- src/hapsira/core/elements.py | 33 +++++++++++++++++++++--------- src/hapsira/twobody/states.py | 6 ++++-- tests/tests_twobody/test_angles.py | 10 +++++++-- 3 files changed, 35 insertions(+), 14 deletions(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 9de13a01a..a9767058b 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -2,7 +2,7 @@ convert between different elements that define the orbit of a body. """ -from math import acos, atan2, cos, log, pi, sin, sqrt, tan +from math import acos, atan, atan2, cos, log, pi, sin, sqrt, tan from numba import njit as jit import numpy as np @@ -38,8 +38,9 @@ "RV2COE_TOL", "rv2coe_hf", "rv2coe_gf", + "mee2coe_hf", + "mee2coe_gf", # TODO - "mee2coe", "mee2rv", ] @@ -515,8 +516,8 @@ def rv2coe_gf(k, r, v, tol, p, ecc, inc, raan, argp, nu): ) -@jit -def mee2coe(p, f, g, h, k, L): +@hjit("Tuple([f,f,f,f,f,f])(f,f,f,f,f,f)") +def mee2coe_hf(p, f, g, h, k, L): r"""Converts from modified equinoctial orbital elements to classical orbital elements. @@ -570,15 +571,27 @@ def mee2coe(p, f, g, h, k, L): arguments. """ - ecc = np.sqrt(f**2 + g**2) - inc = 2 * np.arctan(np.sqrt(h**2 + k**2)) - lonper = np.arctan2(g, f) - raan = np.arctan2(k, h) % (2 * np.pi) - argp = (lonper - raan) % (2 * np.pi) - nu = (L - lonper) % (2 * np.pi) + ecc = sqrt(f**2 + g**2) + inc = 2 * atan(sqrt(h**2 + k**2)) + lonper = atan2(g, f) + raan = atan2(k, h) % (2 * pi) + argp = (lonper - raan) % (2 * pi) + nu = (L - lonper) % (2 * pi) return p, ecc, inc, raan, argp, nu +@gjit( + "void(f,f,f,f,f,f,f[:],f[:],f[:],f[:],f[:],f[:])", + "(),(),(),(),(),()->(),(),(),(),(),()", +) +def mee2coe_gf(p, f, g, h, k, L, p_, ecc, inc, raan, argp, nu): + """ + Vectorized mee2coe + """ + + p_[0], ecc[0], inc[0], raan[0], argp[0], nu[0] = mee2coe_hf(p, f, g, h, k, L) + + @jit def mee2rv(p, f, g, h, k, L): """Calculates position and velocity vector from modified equinoctial elements. diff --git a/src/hapsira/twobody/states.py b/src/hapsira/twobody/states.py index 3b905f418..dda6048f8 100644 --- a/src/hapsira/twobody/states.py +++ b/src/hapsira/twobody/states.py @@ -6,7 +6,7 @@ from hapsira.core.elements import ( coe2mee_gf, coe2rv_gf, - mee2coe, + mee2coe_gf, mee2rv, rv2coe_gf, RV2COE_TOL, @@ -332,7 +332,9 @@ def to_value(self): def to_classical(self): """Converts to classical orbital elements representation.""" - p, ecc, inc, raan, argp, nu = mee2coe(*self.to_value()) + p, ecc, inc, raan, argp, nu = mee2coe_gf( # pylint: disable=E1120,E0633 + *self.to_value() + ) return ClassicalState( self.attractor, diff --git a/tests/tests_twobody/test_angles.py b/tests/tests_twobody/test_angles.py index 453959c23..81ec43309 100644 --- a/tests/tests_twobody/test_angles.py +++ b/tests/tests_twobody/test_angles.py @@ -5,7 +5,13 @@ import pytest from hapsira.bodies import Earth -from hapsira.core.elements import coe2mee_gf, coe2rv_gf, mee2coe, rv2coe_gf, RV2COE_TOL +from hapsira.core.elements import ( + coe2mee_gf, + coe2rv_gf, + mee2coe_gf, + rv2coe_gf, + RV2COE_TOL, +) from hapsira.twobody.angles import ( E_to_M, E_to_nu, @@ -222,7 +228,7 @@ def test_convert_between_coe_and_rv_is_transitive(classical): def test_convert_between_coe_and_mee_is_transitive(classical): - res = mee2coe(*coe2mee_gf(*classical)) # pylint: disable=E1133 + res = mee2coe_gf(*coe2mee_gf(*classical)) # pylint: disable=E1133 assert_allclose(res, classical) From 26044b73d77566ba790a58c50898e30efef7050d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 3 Jan 2024 09:10:35 +0100 Subject: [PATCH 060/346] jit mee2rv --- src/hapsira/core/elements.py | 62 +++++++++++++++++++---------------- src/hapsira/twobody/states.py | 6 ++-- 2 files changed, 38 insertions(+), 30 deletions(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index a9767058b..0ec7c0e8b 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -4,12 +4,7 @@ from math import acos, atan, atan2, cos, log, pi, sin, sqrt, tan -from numba import njit as jit -import numpy as np - -from hapsira.core.angles import E_to_nu_hf, F_to_nu_hf -from hapsira.core.util import rotation_matrix_hf - +from .angles import E_to_nu_hf, F_to_nu_hf from .jit import array_to_V_hf, hjit, gjit, vjit from .math.linalg import ( cross_VV_hf, @@ -22,6 +17,7 @@ sub_VV_hf, transpose_M_hf, ) +from .util import rotation_matrix_hf __all__ = [ @@ -40,8 +36,8 @@ "rv2coe_gf", "mee2coe_hf", "mee2coe_gf", - # TODO - "mee2rv", + "mee2rv_hf", + "mee2rv_gf", ] @@ -126,9 +122,9 @@ def rv_pqw_hf(k, p, ecc, nu): Returns ------- - r: numpy.ndarray + r: tuple[float,float,float] Position. Dimension 3 vector - v: numpy.ndarray + v: tuple[float,float,float] Velocity. Dimension 3 vector Notes @@ -592,8 +588,8 @@ def mee2coe_gf(p, f, g, h, k, L, p_, ecc, inc, raan, argp, nu): p_[0], ecc[0], inc[0], raan[0], argp[0], nu[0] = mee2coe_hf(p, f, g, h, k, L) -@jit -def mee2rv(p, f, g, h, k, L): +@hjit("Tuple([V,V])(f,f,f,f,f,f)") +def mee2rv_hf(p, f, g, h, k, L): # TODO untested """Calculates position and velocity vector from modified equinoctial elements. Parameters @@ -613,9 +609,9 @@ def mee2rv(p, f, g, h, k, L): Returns ------- - r: numpy.ndarray + r: tuple[float,float,float] Position vector. - v: numpy.ndarray + v: tuple[float,float,float] Velocity vector. Note @@ -625,22 +621,22 @@ def mee2rv(p, f, g, h, k, L): Equation 3a and 3b. """ - w = 1 + f * np.cos(L) + g * np.sin(L) + w = 1 + f * cos(L) + g * sin(L) r = p / w s2 = 1 + h**2 + k**2 alpha2 = h**2 - k**2 - rx = (r / s2)(np.cos(L) + alpha2**2 * np.cos(L) + 2 * h * k * np.sin(L)) - ry = (r / s2)(np.sin(L) - alpha2**2 * np.sin(L) + 2 * h * k * np.cos(L)) - rz = (2 * r / s2)(h * np.sin(L) - k * np.cos(L)) + rx = (r / s2) * (cos(L) + alpha2**2 * cos(L) + 2 * h * k * sin(L)) + ry = (r / s2) * (sin(L) - alpha2**2 * sin(L) + 2 * h * k * cos(L)) + rz = (2 * r / s2) * (h * sin(L) - k * cos(L)) vx = ( (-1 / s2) - * (np.sqrt(k / p)) + * (sqrt(k / p)) * ( - np.sin(L) - + alpha2 * np.sin(L) - - 2 * h * k * np.cos(L) + sin(L) + + alpha2 * sin(L) + - 2 * h * k * cos(L) + g - 2 * f * h * k + alpha2 * g @@ -648,16 +644,26 @@ def mee2rv(p, f, g, h, k, L): ) vy = ( (-1 / s2) - * (np.sqrt(k / p)) + * (sqrt(k / p)) * ( - -np.cos(L) - + alpha2 * np.cos(L) - + 2 * h * k * np.sin(L) + -cos(L) + + alpha2 * cos(L) + + 2 * h * k * sin(L) - f + 2 * g * h * k + alpha2 * f ) ) - vz = (2 / s2) * (np.sqrt(k / p)) * (h * np.cos(L) + k * np.sin(L) + f * h + g * k) + vz = (2 / s2) * (sqrt(k / p)) * (h * cos(L) + k * sin(L) + f * h + g * k) + + return (rx, ry, rz), (vx, vy, vz) + + +@gjit("void(f,f,f,f,f,f,u1[:],f[:],f[:])", "(),(),(),(),(),(),(n)->(n),(n)") +def mee2rv_gf(p, f, g, h, k, L, dummy, r, v): + """ + Vectorized mee2rv + """ + assert dummy.shape == (3,) - return np.array([rx, ry, rz]), np.array([vx, vy, vz]) + (r[0], r[1], r[2]), (v[0], v[1], v[2]) = mee2rv_hf(p, f, g, h, k, L) diff --git a/src/hapsira/twobody/states.py b/src/hapsira/twobody/states.py index dda6048f8..1462cf4c1 100644 --- a/src/hapsira/twobody/states.py +++ b/src/hapsira/twobody/states.py @@ -7,7 +7,7 @@ coe2mee_gf, coe2rv_gf, mee2coe_gf, - mee2rv, + mee2rv_gf, rv2coe_gf, RV2COE_TOL, ) @@ -351,5 +351,7 @@ def to_classical(self): def to_vectors(self): """Converts to position and velocity vector representation.""" - r, v = mee2rv(*self.to_value()) + r, v = mee2rv_gf( # pylint: disable=E1120,E0633 + *self.to_value(), np.zeros((3,), dtype="u1") + ) return RVState(self.attractor, (r << u.km, v << u.km / u.s), self.plane) From 76042485ee7d49559146fbb463dce868360430b8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 6 Jan 2024 14:42:54 +0100 Subject: [PATCH 061/346] enable caching --- src/hapsira/core/jit.py | 4 ++++ src/hapsira/settings.py | 6 ++++++ 2 files changed, 10 insertions(+) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index a98d4ffd1..2abdc6bac 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -91,12 +91,14 @@ def wrapper(inner_func: Callable) -> Callable: cfg = dict( device=True, inline=settings["INLINE"].value, + cache=settings["CACHE"].value, ) else: wjit = nb.jit cfg = dict( nopython=settings["NOPYTHON"].value, inline="always" if settings["INLINE"].value else "never", + cache=settings["CACHE"].value, ) cfg.update(kwargs) @@ -140,6 +142,7 @@ def wrapper(inner_func: Callable) -> Callable: cfg = dict( target=settings["TARGET"].value, + cache=settings["CACHE"].value, ) if settings["TARGET"].value != "cuda": cfg["nopython"] = settings["NOPYTHON"].value @@ -185,6 +188,7 @@ def wrapper(inner_func: Callable) -> Callable: cfg = dict( target=settings["TARGET"].value, + cache=settings["CACHE"].value, ) if settings["TARGET"].value != "cuda": cfg["nopython"] = settings["NOPYTHON"].value diff --git a/src/hapsira/settings.py b/src/hapsira/settings.py index c45cdbf79..f12a4a881 100644 --- a/src/hapsira/settings.py +++ b/src/hapsira/settings.py @@ -101,6 +101,12 @@ def __init__(self): False, ) ) + self._add( + Setting( + "CACHE", + not self["DEBUG"].value, + ) + ) self._add( Setting( "LOGLEVEL", From 07db36939c69965e0adf66dfd313c3e82755f380 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 6 Jan 2024 14:43:46 +0100 Subject: [PATCH 062/346] jit iod --- src/hapsira/core/iod.py | 631 ++++++++++++++++++-------------- src/hapsira/core/math/linalg.py | 5 + src/hapsira/iod/izzo.py | 7 +- src/hapsira/iod/vallado.py | 7 +- tests/test_iod.py | 12 +- 5 files changed, 379 insertions(+), 283 deletions(-) diff --git a/src/hapsira/core/iod.py b/src/hapsira/core/iod.py index 3b07140ec..1fef127ce 100644 --- a/src/hapsira/core/iod.py +++ b/src/hapsira/core/iod.py @@ -1,14 +1,282 @@ -from numba import njit as jit -import numpy as np -from numpy import cross, pi - -from .jit import array_to_V_hf -from .math.linalg import norm_hf +from math import acos, asinh, exp, floor, inf, log, pi, sqrt + +from .jit import array_to_V_hf, hjit, gjit, vjit +from .math.linalg import ( + add_VV_hf, + cross_VV_hf, + div_Vs_hf, + matmul_VV_hf, + mul_Vs_hf, + norm_hf, + sub_VV_hf, +) from .math.special import hyp2f1b_hf, stumpff_c2_hf, stumpff_c3_hf -@jit -def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): +__all__ = [ + "vallado_hf", + "vallado_gf", + "izzo_hf", + "izzo_gf", + "compute_T_min_hf", + "compute_T_min_gf", + "compute_y_hf", + "compute_y_vf", + "tof_equation_y_hf", + "tof_equation_y_vf", +] + + +@hjit("f(f,f)") +def compute_y_hf(x, ll): + """Computes y.""" + return sqrt(1 - ll**2 * (1 - x**2)) + + +@vjit("f(f,f)") +def compute_y_vf(x, ll): + """ + Vectorized compute_y + """ + + return compute_y_hf(x, ll) + + +@hjit("f(f,f,f,f)") +def _tof_equation_p_hf(x, y, T, ll): + # TODO: What about derivatives when x approaches 1? + return (3 * T * x - 2 + 2 * ll**3 * x / y) / (1 - x**2) + + +@hjit("f(f,f,f,f,f)") +def _tof_equation_p2_hf(x, y, T, dT, ll): + return (3 * T + 5 * x * dT + 2 * (1 - ll**2) * ll**3 / y**3) / (1 - x**2) + + +@hjit("f(f,f,f,f,f,f)") +def _tof_equation_p3_hf(x, y, _, dT, ddT, ll): + return (7 * x * ddT + 8 * dT - 6 * (1 - ll**2) * ll**5 * x / y**5) / ( + 1 - x**2 + ) + + +@hjit("f(f,f,i8,b1)") +def _initial_guess_hf(T, ll, M, lowpath): + """Initial guess.""" + if M == 0: + # Single revolution + T_0 = acos(ll) + ll * sqrt(1 - ll**2) + M * pi # Equation 19 + T_1 = 2 * (1 - ll**3) / 3 # Equation 21 + if T >= T_0: + x_0 = (T_0 / T) ** (2 / 3) - 1 + elif T < T_1: + x_0 = 5 / 2 * T_1 / T * (T_1 - T) / (1 - ll**5) + 1 + else: + # This is the real condition, which is not exactly equivalent + # elif T_1 < T < T_0 + # Corrected initial guess, + # piecewise equation right after expression (30) in the original paper is incorrect + # See https://github.com/poliastro/poliastro/issues/1362 + x_0 = exp(log(2) * log(T / T_0) / log(T_1 / T_0)) - 1 + + return x_0 + else: + # Multiple revolution + x_0l = (((M * pi + pi) / (8 * T)) ** (2 / 3) - 1) / ( + ((M * pi + pi) / (8 * T)) ** (2 / 3) + 1 + ) + x_0r = (((8 * T) / (M * pi)) ** (2 / 3) - 1) / ( + ((8 * T) / (M * pi)) ** (2 / 3) + 1 + ) + + # Select one of the solutions according to desired type of path + x_0 = max((x_0l, x_0r)) if lowpath else min((x_0l, x_0r)) + + return x_0 + + +@hjit("Tuple([f,f,f,f])(f,f,f,f,f,f,f,f)") +def _reconstruct_hf(x, y, r1, r2, ll, gamma, rho, sigma): + """Reconstruct solution velocity vectors.""" + V_r1 = gamma * ((ll * y - x) - rho * (ll * y + x)) / r1 + V_r2 = -gamma * ((ll * y - x) + rho * (ll * y + x)) / r2 + V_t1 = gamma * sigma * (y + ll * x) / r1 + V_t2 = gamma * sigma * (y + ll * x) / r2 + return V_r1, V_r2, V_t1, V_t2 + + +@hjit("f(f,f,f)") +def _compute_psi_hf(x, y, ll): + """Computes psi. + + "The auxiliary angle psi is computed using Eq.(17) by the appropriate + inverse function" + + """ + if -1 <= x < 1: + # Elliptic motion + # Use arc cosine to avoid numerical errors + return acos(x * y + ll * (1 - x**2)) + elif x > 1: + # Hyperbolic motion + # The hyperbolic sine is bijective + return asinh((y - x * ll) * sqrt(x**2 - 1)) + else: + # Parabolic motion + return 0.0 + + +@hjit("f(f,f,f,f,i8)") +def tof_equation_y_hf(x, y, T0, ll, M): + """Time of flight equation with externally computated y.""" + if M == 0 and sqrt(0.6) < x < sqrt(1.4): + eta = y - ll * x + S_1 = (1 - ll - x * eta) * 0.5 + Q = 4 / 3 * hyp2f1b_hf(S_1) + T_ = (eta**3 * Q + 4 * ll * eta) * 0.5 + else: + psi = _compute_psi_hf(x, y, ll) + T_ = (((psi + M * pi) / sqrt(abs(1 - x**2))) - x + ll * y) / (1 - x**2) + + return T_ - T0 + + +@vjit("f(f,f,f,f,i8)") +def tof_equation_y_vf(x, y, T0, ll, M): + """ + Vectorized tof_equation_y + """ + + return tof_equation_y_hf(x, y, T0, ll, M) + + +@hjit("f(f,f,f,i8)") +def _tof_equation_hf(x, T0, ll, M): + """Time of flight equation.""" + return tof_equation_y_hf(x, compute_y_hf(x, ll), T0, ll, M) + + +@hjit("f(f,f,f,f,i8)") +def _halley_hf(p0, T0, ll, tol, maxiter): + """Find a minimum of time of flight equation using the Halley method. + + Notes + ----- + This function is private because it assumes a calling convention specific to + this module and is not really reusable. + + """ + for ii in range(maxiter): + y = compute_y_hf(p0, ll) + fder = _tof_equation_p_hf(p0, y, T0, ll) + fder2 = _tof_equation_p2_hf(p0, y, T0, fder, ll) + if fder2 == 0: + raise RuntimeError("Derivative was zero") + fder3 = _tof_equation_p3_hf(p0, y, T0, fder, fder2, ll) + + # Halley step (cubic) + p = p0 - 2 * fder * fder2 / (2 * fder2**2 - fder * fder3) + + if abs(p - p0) < tol: + return p + p0 = p + + raise RuntimeError("Failed to converge") + + +@hjit("Tuple([f,f])(f,i8,i8,f)") +def compute_T_min_hf(ll, M, numiter, rtol): + """Compute minimum T.""" + if ll == 1: + x_T_min = 0.0 + T_min = _tof_equation_hf(x_T_min, 0.0, ll, M) + else: + if M == 0: + x_T_min = inf + T_min = 0.0 + else: + # Set x_i > 0 to avoid problems at ll = -1 + x_i = 0.1 + T_i = _tof_equation_hf(x_i, 0.0, ll, M) + x_T_min = _halley_hf(x_i, T_i, ll, rtol, numiter) + T_min = _tof_equation_hf(x_T_min, 0.0, ll, M) + + return x_T_min, T_min + + +@gjit("void(f,i8,i8,f,f[:],f[:])", "(),(),(),()->(),()") +def compute_T_min_gf(ll, M, numiter, rtol, x_T_min, T_min): + """ + Vectorized compute_T_min + """ + + x_T_min[0], T_min[0] = compute_T_min_hf(ll, M, numiter, rtol) + + +@hjit("f(f,f,f,i8,f,i8)") +def _householder_hf(p0, T0, ll, M, tol, maxiter): + """Find a zero of time of flight equation using the Householder method. + + Notes + ----- + This function is private because it assumes a calling convention specific to + this module and is not really reusable. + + """ + for ii in range(maxiter): + y = compute_y_hf(p0, ll) + fval = tof_equation_y_hf(p0, y, T0, ll, M) + T = fval + T0 + fder = _tof_equation_p_hf(p0, y, T, ll) + fder2 = _tof_equation_p2_hf(p0, y, T, fder, ll) + fder3 = _tof_equation_p3_hf(p0, y, T, fder, fder2, ll) + + # Householder step (quartic) + p = p0 - fval * ( + (fder**2 - fval * fder2 / 2) + / (fder * (fder**2 - fval * fder2) + fder3 * fval**2 / 6) + ) + + if abs(p - p0) < tol: + return p + p0 = p + + raise RuntimeError("Failed to converge") + + +@hjit("Tuple([f,f])(f,f,i8,i8,b1,f)") +def _find_xy_hf(ll, T, M, numiter, lowpath, rtol): + """Computes all x, y for given number of revolutions.""" + # For abs(ll) == 1 the derivative is not continuous + assert abs(ll) < 1 + assert T > 0 # Mistake in the original paper + + M_max = floor(T / pi) + T_00 = acos(ll) + ll * sqrt(1 - ll**2) # T_xM + + # Refine maximum number of revolutions if necessary + if T < T_00 + M_max * pi and M_max > 0: + _, T_min = compute_T_min_hf(ll, M_max, numiter, rtol) + if T < T_min: + M_max -= 1 + + # Check if a feasible solution exist for the given number of revolutions + # This departs from the original paper in that we do not compute all solutions + if M > M_max: + raise ValueError("No feasible solution, try lower M") + + # Initial guess + x_0 = _initial_guess_hf(T, ll, M, lowpath) + + # Start Householder iterations from x_0 and find x, y + x = _householder_hf(x_0, T, ll, M, rtol, numiter) + y = compute_y_hf(x, ll) + + return x, y + + +@hjit("Tuple([V,V])(f,V,V,f,i8,b1,b1,i8,f)") +def vallado_hf(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): r"""Solves the Lambert's problem. The algorithm returns the initial velocity vector and the final one, these are @@ -54,9 +322,9 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): ---------- k : float Gravitational Parameter - r0 : numpy.ndarray + r0 : tuple[float,float,float] Initial position vector - r : numpy.ndarray + r : tuple[float,float,float] Final position vector tof : float Time of flight @@ -74,9 +342,9 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): Returns ------- - v0: numpy.ndarray + v0: tuple[float,float,float] Initial velocity vector - v: numpy.ndarray + v: tuple[float,float,float] Final velocity vector Examples @@ -103,20 +371,20 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): """ # TODO: expand for the multi-revolution case. - # Issue: https://github.com/hapsira/hapsira/issues/858 + # Issue: https://github.com/poliastro/poliastro/issues/858 if M > 0: raise NotImplementedError( - "Multi-revolution scenario not supported for Vallado. See issue https://github.com/hapsira/hapsira/issues/858" + "Multi-revolution scenario not supported for Vallado. See issue https://github.com/poliastro/poliastro/issues/858" ) t_m = 1 if prograde else -1 - norm_r0 = norm_hf(array_to_V_hf(r0)) - norm_r = norm_hf(array_to_V_hf(r)) + norm_r0 = norm_hf(r0) + norm_r = norm_hf(r) norm_r0_times_norm_r = norm_r0 * norm_r norm_r0_plus_norm_r = norm_r0 + norm_r - cos_dnu = (r0 @ r) / norm_r0_times_norm_r + cos_dnu = matmul_VV_hf(r0, r) / norm_r0_times_norm_r A = t_m * (norm_r * norm_r0 * (1 + cos_dnu)) ** 0.5 @@ -124,8 +392,8 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): raise RuntimeError("Cannot compute orbit, phase angle is 180 degrees") psi = 0.0 - psi_low = -4 * np.pi**2 - psi_up = 4 * np.pi**2 + psi_low = -4 * pi**2 + psi_up = 4 * pi**2 count = 0 @@ -142,18 +410,18 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): psi = ( 0.8 * (1.0 / stumpff_c3_hf(psi)) - * (1.0 - norm_r0_times_norm_r * np.sqrt(stumpff_c2_hf(psi)) / A) + * (1.0 - norm_r0_times_norm_r * sqrt(stumpff_c2_hf(psi)) / A) ) y = ( norm_r0_plus_norm_r + A * (psi * stumpff_c3_hf(psi) - 1) / stumpff_c2_hf(psi) ** 0.5 ) - xi = np.sqrt(y / stumpff_c2_hf(psi)) - tof_new = (xi**3 * stumpff_c3_hf(psi) + A * np.sqrt(y)) / np.sqrt(k) + xi = sqrt(y / stumpff_c2_hf(psi)) + tof_new = (xi**3 * stumpff_c3_hf(psi) + A * sqrt(y)) / sqrt(k) # Convergence check - if np.abs((tof_new - tof) / tof) < rtol: + if abs((tof_new - tof) / tof) < rtol: break count += 1 # Bisection check @@ -166,27 +434,41 @@ def vallado(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): raise RuntimeError("Maximum number of iterations reached") f = 1 - y / norm_r0 - g = A * np.sqrt(y / k) + g = A * sqrt(y / k) gdot = 1 - y / norm_r - v0 = (r - f * r0) / g - v = (gdot * r - r0) / g + v0 = div_Vs_hf(sub_VV_hf(r, mul_Vs_hf(r0, f)), g) + v = div_Vs_hf(sub_VV_hf(mul_Vs_hf(r, gdot), r0), g) return v0, v -@jit -def izzo(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol): +@gjit( + "void(f,f[:],f[:],f,i8,b1,b1,i8,f,f[:],f[:])", + "(),(n),(n),(),(),(),(),(),()->(n),(n)", +) +def vallado_gf(k, r0, r, tof, M, prograde, lowpath, numiter, rtol, v0, v): + """ + Vectorized vallado + """ + + ((v0[0], v0[1], v0[2]), (v[0], v[1], v[2])) = vallado_hf( + k, array_to_V_hf(r0), array_to_V_hf(r), tof, M, prograde, lowpath, numiter, rtol + ) + + +@hjit("Tuple([V,V])(f,V,V,f,i8,b1,b1,i8,f)") +def izzo_hf(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol): """Aplies izzo algorithm to solve Lambert's problem. Parameters ---------- k : float Gravitational Constant - r1 : numpy.ndarray + r1 : tuple[float,float,float] Initial position vector - r2 : numpy.ndarray + r2 : tuple[float,float,float] Final position vector tof : float Time of flight between both positions @@ -204,9 +486,9 @@ def izzo(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol): Returns ------- - v1: numpy.ndarray + v1: tuple[float,float,float] Initial velocity vector - v2: numpy.ndarray + v2: tuple[float,float,float] Final velocity vector """ @@ -215,282 +497,83 @@ def izzo(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol): assert k > 0 # Check collinearity of r1 and r2 - if not cross(r1, r2).any(): + cl = cross_VV_hf(r1, r2) + if cl[0] == 0 and cl[1] == 0 and cl[2] == 0: raise ValueError("Lambert solution cannot be computed for collinear vectors") # Chord - c = r2 - r1 + c = sub_VV_hf(r2, r1) c_norm, r1_norm, r2_norm = ( - norm_hf(array_to_V_hf(c)), - norm_hf(array_to_V_hf(r1)), - norm_hf(array_to_V_hf(r2)), + norm_hf(c), + norm_hf(r1), + norm_hf(r2), ) # Semiperimeter s = (r1_norm + r2_norm + c_norm) * 0.5 # Versors - i_r1, i_r2 = r1 / r1_norm, r2 / r2_norm - i_h = cross(i_r1, i_r2) - i_h = i_h / norm_hf(array_to_V_hf(i_h)) # Fixed from paper + i_r1, i_r2 = div_Vs_hf(r1, r1_norm), div_Vs_hf(r2, r2_norm) + i_h = cross_VV_hf(i_r1, i_r2) + i_h = div_Vs_hf(i_h, norm_hf(i_h)) # Fixed from paper # Geometry of the problem - ll = np.sqrt(1 - min(1.0, c_norm / s)) + ll = sqrt(1 - min(1.0, c_norm / s)) # Compute the fundamental tangential directions if i_h[2] < 0: ll = -ll - i_t1, i_t2 = cross(i_r1, i_h), cross(i_r2, i_h) + i_t1, i_t2 = cross_VV_hf(i_r1, i_h), cross_VV_hf(i_r2, i_h) else: - i_t1, i_t2 = cross(i_h, i_r1), cross(i_h, i_r2) + i_t1, i_t2 = cross_VV_hf(i_h, i_r1), cross_VV_hf(i_h, i_r2) # Correct transfer angle parameter and tangential vectors if required - ll, i_t1, i_t2 = (ll, i_t1, i_t2) if prograde else (-ll, -i_t1, -i_t2) + ll, i_t1, i_t2 = ( + (ll, i_t1, i_t2) + if prograde + else (-ll, mul_Vs_hf(i_t1, -1), mul_Vs_hf(i_t2, -1)) + ) # Non dimensional time of flight - T = np.sqrt(2 * k / s**3) * tof + T = sqrt(2 * k / s**3) * tof # Find solutions - x, y = _find_xy(ll, T, M, numiter, lowpath, rtol) + x, y = _find_xy_hf(ll, T, M, numiter, lowpath, rtol) # Reconstruct - gamma = np.sqrt(k * s / 2) + gamma = sqrt(k * s / 2) rho = (r1_norm - r2_norm) / c_norm - sigma = np.sqrt(1 - rho**2) + sigma = sqrt(1 - rho**2) # Compute the radial and tangential components at r0 and r - V_r1, V_r2, V_t1, V_t2 = _reconstruct(x, y, r1_norm, r2_norm, ll, gamma, rho, sigma) + V_r1, V_r2, V_t1, V_t2 = _reconstruct_hf( + x, y, r1_norm, r2_norm, ll, gamma, rho, sigma + ) # Solve for the initial and final velocity - v1 = V_r1 * (r1 / r1_norm) + V_t1 * i_t1 - v2 = V_r2 * (r2 / r2_norm) + V_t2 * i_t2 + v1 = add_VV_hf(mul_Vs_hf(div_Vs_hf(r1, r1_norm), V_r1), mul_Vs_hf(i_t1, V_t1)) + v2 = add_VV_hf(mul_Vs_hf(div_Vs_hf(r2, r2_norm), V_r2), mul_Vs_hf(i_t2, V_t2)) return v1, v2 -@jit -def _reconstruct(x, y, r1, r2, ll, gamma, rho, sigma): - """Reconstruct solution velocity vectors.""" - V_r1 = gamma * ((ll * y - x) - rho * (ll * y + x)) / r1 - V_r2 = -gamma * ((ll * y - x) + rho * (ll * y + x)) / r2 - V_t1 = gamma * sigma * (y + ll * x) / r1 - V_t2 = gamma * sigma * (y + ll * x) / r2 - return V_r1, V_r2, V_t1, V_t2 - - -@jit -def _find_xy(ll, T, M, numiter, lowpath, rtol): - """Computes all x, y for given number of revolutions.""" - # For abs(ll) == 1 the derivative is not continuous - assert abs(ll) < 1 - assert T > 0 # Mistake in the original paper - - M_max = np.floor(T / pi) - T_00 = np.arccos(ll) + ll * np.sqrt(1 - ll**2) # T_xM - - # Refine maximum number of revolutions if necessary - if T < T_00 + M_max * pi and M_max > 0: - _, T_min = _compute_T_min(ll, M_max, numiter, rtol) - if T < T_min: - M_max -= 1 - - # Check if a feasible solution exist for the given number of revolutions - # This departs from the original paper in that we do not compute all solutions - if M > M_max: - raise ValueError("No feasible solution, try lower M") - - # Initial guess - x_0 = _initial_guess(T, ll, M, lowpath) - - # Start Householder iterations from x_0 and find x, y - x = _householder(x_0, T, ll, M, rtol, numiter) - y = _compute_y(x, ll) - - return x, y - - -@jit -def _compute_y(x, ll): - """Computes y.""" - return np.sqrt(1 - ll**2 * (1 - x**2)) - - -@jit -def _compute_psi(x, y, ll): - """Computes psi. - - "The auxiliary angle psi is computed using Eq.(17) by the appropriate - inverse function" - - """ - if -1 <= x < 1: - # Elliptic motion - # Use arc cosine to avoid numerical errors - return np.arccos(x * y + ll * (1 - x**2)) - elif x > 1: - # Hyperbolic motion - # The hyperbolic sine is bijective - return np.arcsinh((y - x * ll) * np.sqrt(x**2 - 1)) - else: - # Parabolic motion - return 0.0 - - -@jit -def _tof_equation(x, T0, ll, M): - """Time of flight equation.""" - return _tof_equation_y(x, _compute_y(x, ll), T0, ll, M) - - -@jit -def _tof_equation_y(x, y, T0, ll, M): - """Time of flight equation with externally computated y.""" - if M == 0 and np.sqrt(0.6) < x < np.sqrt(1.4): - eta = y - ll * x - S_1 = (1 - ll - x * eta) * 0.5 - Q = 4 / 3 * hyp2f1b_hf(S_1) - T_ = (eta**3 * Q + 4 * ll * eta) * 0.5 - else: - psi = _compute_psi(x, y, ll) - T_ = np.divide( - np.divide(psi + M * pi, np.sqrt(np.abs(1 - x**2))) - x + ll * y, - (1 - x**2), - ) - - return T_ - T0 - - -@jit -def _tof_equation_p(x, y, T, ll): - # TODO: What about derivatives when x approaches 1? - return (3 * T * x - 2 + 2 * ll**3 * x / y) / (1 - x**2) - - -@jit -def _tof_equation_p2(x, y, T, dT, ll): - return (3 * T + 5 * x * dT + 2 * (1 - ll**2) * ll**3 / y**3) / (1 - x**2) - - -@jit -def _tof_equation_p3(x, y, _, dT, ddT, ll): - return (7 * x * ddT + 8 * dT - 6 * (1 - ll**2) * ll**5 * x / y**5) / ( - 1 - x**2 - ) - - -@jit -def _compute_T_min(ll, M, numiter, rtol): - """Compute minimum T.""" - if ll == 1: - x_T_min = 0.0 - T_min = _tof_equation(x_T_min, 0.0, ll, M) - else: - if M == 0: - x_T_min = np.inf - T_min = 0.0 - else: - # Set x_i > 0 to avoid problems at ll = -1 - x_i = 0.1 - T_i = _tof_equation(x_i, 0.0, ll, M) - x_T_min = _halley(x_i, T_i, ll, rtol, numiter) - T_min = _tof_equation(x_T_min, 0.0, ll, M) - - return x_T_min, T_min - - -@jit -def _initial_guess(T, ll, M, lowpath): - """Initial guess.""" - if M == 0: - # Single revolution - T_0 = np.arccos(ll) + ll * np.sqrt(1 - ll**2) + M * pi # Equation 19 - T_1 = 2 * (1 - ll**3) / 3 # Equation 21 - if T >= T_0: - x_0 = (T_0 / T) ** (2 / 3) - 1 - elif T < T_1: - x_0 = 5 / 2 * T_1 / T * (T_1 - T) / (1 - ll**5) + 1 - else: - # This is the real condition, which is not exactly equivalent - # elif T_1 < T < T_0 - # Corrected initial guess, - # piecewise equation right after expression (30) in the original paper is incorrect - # See https://github.com/hapsira/hapsira/issues/1362 - x_0 = np.exp(np.log(2) * np.log(T / T_0) / np.log(T_1 / T_0)) - 1 - - return x_0 - else: - # Multiple revolution - x_0l = (((M * pi + pi) / (8 * T)) ** (2 / 3) - 1) / ( - ((M * pi + pi) / (8 * T)) ** (2 / 3) + 1 - ) - x_0r = (((8 * T) / (M * pi)) ** (2 / 3) - 1) / ( - ((8 * T) / (M * pi)) ** (2 / 3) + 1 - ) - - # Select one of the solutions according to desired type of path - x_0 = ( - np.max(np.array([x_0l, x_0r])) - if lowpath - else np.min(np.array([x_0l, x_0r])) - ) - - return x_0 - - -@jit -def _halley(p0, T0, ll, tol, maxiter): - """Find a minimum of time of flight equation using the Halley method. - - Notes - ----- - This function is private because it assumes a calling convention specific to - this module and is not really reusable. - +@gjit( + "void(f,f[:],f[:],f,i8,b1,b1,i8,f,f[:],f[:])", + "(),(n),(n),(),(),(),(),(),()->(n),(n)", +) +def izzo_gf(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol, v1, v2): """ - for ii in range(maxiter): - y = _compute_y(p0, ll) - fder = _tof_equation_p(p0, y, T0, ll) - fder2 = _tof_equation_p2(p0, y, T0, fder, ll) - if fder2 == 0: - raise RuntimeError("Derivative was zero") - fder3 = _tof_equation_p3(p0, y, T0, fder, fder2, ll) - - # Halley step (cubic) - p = p0 - 2 * fder * fder2 / (2 * fder2**2 - fder * fder3) - - if abs(p - p0) < tol: - return p - p0 = p - - raise RuntimeError("Failed to converge") - - -@jit -def _householder(p0, T0, ll, M, tol, maxiter): - """Find a zero of time of flight equation using the Householder method. - - Notes - ----- - This function is private because it assumes a calling convention specific to - this module and is not really reusable. - + Vectorized izzo """ - for ii in range(maxiter): - y = _compute_y(p0, ll) - fval = _tof_equation_y(p0, y, T0, ll, M) - T = fval + T0 - fder = _tof_equation_p(p0, y, T, ll) - fder2 = _tof_equation_p2(p0, y, T, fder, ll) - fder3 = _tof_equation_p3(p0, y, T, fder, fder2, ll) - - # Householder step (quartic) - p = p0 - fval * ( - (fder**2 - fval * fder2 / 2) - / (fder * (fder**2 - fval * fder2) + fder3 * fval**2 / 6) - ) - if abs(p - p0) < tol: - return p - p0 = p - - raise RuntimeError("Failed to converge") + ((v1[0], v1[1], v1[2]), (v2[0], v2[1], v2[2])) = izzo_hf( + k, + array_to_V_hf(r1), + array_to_V_hf(r2), + tof, + M, + prograde, + lowpath, + numiter, + rtol, + ) diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 30f0ce1fc..7f42604a7 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -16,6 +16,11 @@ ] +@hjit("V(V,V)") +def add_VV_hf(a, b): + return a[0] + b[0], a[1] + b[1], a[2] + b[2] + + @hjit("V(V,V)") def cross_VV_hf(a, b): return ( diff --git a/src/hapsira/iod/izzo.py b/src/hapsira/iod/izzo.py index 42af60b85..8832459ed 100644 --- a/src/hapsira/iod/izzo.py +++ b/src/hapsira/iod/izzo.py @@ -1,7 +1,8 @@ """Izzo's algorithm for Lambert's problem.""" from astropy import units as u +import numpy as np -from hapsira.core.iod import izzo as izzo_fast +from hapsira.core.iod import izzo_gf kms = u.km / u.s @@ -44,5 +45,7 @@ def lambert(k, r0, r, tof, M=0, prograde=True, lowpath=True, numiter=35, rtol=1e r_ = r.to_value(u.km) tof_ = tof.to_value(u.s) - v0, v = izzo_fast(k_, r0_, r_, tof_, M, prograde, lowpath, numiter, rtol) + v0, v = izzo_gf( # pylint: disable=E1120,E0633 + k_, r0_, r_, tof_, M, np.asarray(prograde), np.asarray(lowpath), numiter, rtol + ) return v0 << kms, v << kms diff --git a/src/hapsira/iod/vallado.py b/src/hapsira/iod/vallado.py index a267d744e..7b2e40b46 100644 --- a/src/hapsira/iod/vallado.py +++ b/src/hapsira/iod/vallado.py @@ -1,7 +1,8 @@ """Initial orbit determination.""" from astropy import units as u +import numpy as np -from hapsira.core.iod import vallado as vallado_fast +from hapsira.core.iod import vallado_gf kms = u.km / u.s @@ -55,6 +56,8 @@ def lambert(k, r0, r, tof, M=0, prograde=True, lowpath=True, numiter=35, rtol=1e r_ = r.to_value(u.km) tof_ = tof.to_value(u.s) - v0, v = vallado_fast(k_, r0_, r_, tof_, M, prograde, lowpath, numiter, rtol) + v0, v = vallado_gf( # pylint: disable=E1120,E0633 + k_, r0_, r_, tof_, M, np.asarray(prograde), np.asarray(lowpath), numiter, rtol + ) return v0 << kms, v << kms diff --git a/tests/test_iod.py b/tests/test_iod.py index ba44a1f37..63aa90f02 100644 --- a/tests/test_iod.py +++ b/tests/test_iod.py @@ -3,7 +3,7 @@ import pytest from hapsira.bodies import Earth -from hapsira.core import iod +from hapsira.core.iod import compute_T_min_gf, compute_y_vf, tof_equation_y_vf from hapsira.iod import izzo, vallado @@ -119,9 +119,11 @@ def test_collinear_vectors_input(lambert): @pytest.mark.parametrize("M", [1, 2, 3]) def test_minimum_time_of_flight_convergence(M): ll = -1 - x_T_min_expected, T_min_expected = iod._compute_T_min(ll, M, numiter=10, rtol=1e-8) - y = iod._compute_y(x_T_min_expected, ll) - T_min = iod._tof_equation_y(x_T_min_expected, y, 0.0, ll, M) + x_T_min_expected, T_min_expected = compute_T_min_gf( # pylint: disable=E1120,E0633 + ll, M, 10, 1e-8 + ) + y = compute_y_vf(x_T_min_expected, ll) + T_min = tof_equation_y_vf(x_T_min_expected, y, 0.0, ll, M) assert T_min_expected == T_min @@ -171,6 +173,6 @@ def test_vallado_not_implemented_multirev(): with pytest.raises(NotImplementedError) as excinfo: vallado.lambert(k, r0, r, tof, M=1) assert ( - "Multi-revolution scenario not supported for Vallado. See issue https://github.com/hapsira/hapsira/issues/858" + "Multi-revolution scenario not supported for Vallado. See issue https://github.com/poliastro/poliastro/issues/858" in excinfo.exconly() ) From 3c9f7aee59433c282ecee7324c435d4cebad7f62 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 6 Jan 2024 15:48:42 +0100 Subject: [PATCH 063/346] jit danby --- src/hapsira/core/math/linalg.py | 11 +++ src/hapsira/core/propagation/__init__.py | 7 +- src/hapsira/core/propagation/danby.py | 112 +++++++++++++++------- src/hapsira/twobody/propagation/danby.py | 6 +- tests/tests_core/test_core_propagation.py | 4 +- 5 files changed, 99 insertions(+), 41 deletions(-) diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 7f42604a7..d7719505e 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -3,6 +3,7 @@ from ..jit import hjit, vjit __all__ = [ + "add_VV_hf", "cross_VV_hf", "div_Vs_hf", "matmul_MM_hf", @@ -11,6 +12,7 @@ "mul_Vs_hf", "norm_hf", "norm_vf", + "sign_hf", "sub_VV_hf", "transpose_M_hf", ] @@ -86,6 +88,15 @@ def norm_vf(a, b, c): return norm_hf((a, b, c)) +@hjit("f(f)") +def sign_hf(x): + if x < 0.0: + return -1.0 + if x == 0.0: + return 0.0 + return 1.0 # if x > 0 + + @hjit("V(V,V)") def sub_VV_hf(va, vb): return va[0] - vb[0], va[1] - vb[1], va[2] - vb[2] diff --git a/src/hapsira/core/propagation/__init__.py b/src/hapsira/core/propagation/__init__.py index 5460386d5..93322cce4 100644 --- a/src/hapsira/core/propagation/__init__.py +++ b/src/hapsira/core/propagation/__init__.py @@ -2,7 +2,8 @@ from hapsira.core.propagation.base import func_twobody from hapsira.core.propagation.cowell import cowell -from hapsira.core.propagation.danby import danby, danby_coe + +# from hapsira.core.propagation.danby import danby, danby_coe from hapsira.core.propagation.farnocchia import ( farnocchia_coe, farnocchia_rv as farnocchia, @@ -28,8 +29,8 @@ "pimienta", "gooding_coe", "gooding", - "danby_coe", - "danby", + # "danby_coe", + # "danby", "recseries_coe", "recseries", ] diff --git a/src/hapsira/core/propagation/danby.py b/src/hapsira/core/propagation/danby.py index b4530797c..d428569b9 100644 --- a/src/hapsira/core/propagation/danby.py +++ b/src/hapsira/core/propagation/danby.py @@ -1,64 +1,82 @@ -from numba import njit as jit -import numpy as np +from math import atan2, cos, cosh, floor, log, pi, sin, sinh, sqrt -from hapsira.core.angles import E_to_M_hf, F_to_M_hf, nu_to_E_hf, nu_to_F_hf -from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL -from ..jit import array_to_V_hf +from ..angles import E_to_M_hf, F_to_M_hf, nu_to_E_hf, nu_to_F_hf +from ..elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..math.linalg import sign_hf +from ..jit import array_to_V_hf, hjit, gjit, vjit -@jit -def danby_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter=20, rtol=1e-8): +__all__ = [ + "danby_coe_hf", + "danby_coe_vf", + "danby_rv_hf", + "danby_rv_gf", + "DANBY_NUMITER", + "DANBY_RTOL", +] + + +DANBY_NUMITER = 20 +DANBY_RTOL = 1e-8 + + +@hjit("f(f,f,f,f,f,f,f,f,i8,f)") +def danby_coe_hf(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol): + """ + Scalar danby_coe + """ + semi_axis_a = p / (1 - ecc**2) - n = np.sqrt(k / np.abs(semi_axis_a) ** 3) + n = sqrt(k / abs(semi_axis_a) ** 3) if ecc == 0: # Solving for circular orbit M0 = nu # for circular orbit M = E = nu M = M0 + n * tof - nu = M - 2 * np.pi * np.floor(M / 2 / np.pi) + nu = M - 2 * pi * floor(M / 2 / pi) return nu elif ecc < 1.0: # For elliptical orbit M0 = E_to_M_hf(nu_to_E_hf(nu, ecc), ecc) M = M0 + n * tof - xma = M - 2 * np.pi * np.floor(M / 2 / np.pi) - E = xma + 0.85 * np.sign(np.sin(xma)) * ecc + xma = M - 2 * pi * floor(M / 2 / pi) + E = xma + 0.85 * sign_hf(sin(xma)) * ecc else: # For parabolic and hyperbolic M0 = F_to_M_hf(nu_to_F_hf(nu, ecc), ecc) M = M0 + n * tof - xma = M - 2 * np.pi * np.floor(M / 2 / np.pi) - E = np.log(2 * xma / ecc + 1.8) + xma = M - 2 * pi * floor(M / 2 / pi) + E = log(2 * xma / ecc + 1.8) # Iterations begin n = 0 while n <= numiter: if ecc < 1.0: - s = ecc * np.sin(E) - c = ecc * np.cos(E) + s = ecc * sin(E) + c = ecc * cos(E) f = E - s - xma fp = 1 - c fpp = s fppp = c else: - s = ecc * np.sinh(E) - c = ecc * np.cosh(E) + s = ecc * sinh(E) + c = ecc * cosh(E) f = s - E - xma fp = c - 1 fpp = s fppp = c - if np.abs(f) <= rtol: + if abs(f) <= rtol: if ecc < 1.0: - sta = np.sqrt(1 - ecc**2) * np.sin(E) - cta = np.cos(E) - ecc + sta = sqrt(1 - ecc**2) * sin(E) + cta = cos(E) - ecc else: - sta = np.sqrt(ecc**2 - 1) * np.sinh(E) - cta = ecc - np.cosh(E) + sta = sqrt(ecc**2 - 1) * sinh(E) + cta = ecc - cosh(E) - nu = np.arctan2(sta, cta) + nu = atan2(sta, cta) break else: delta = -f / fp @@ -72,8 +90,17 @@ def danby_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter=20, rtol=1e-8): return nu -@jit -def danby(k, r0, v0, tof, numiter=20, rtol=1e-8): +@vjit("f(f,f,f,f,f,f,f,f,i8,f)") +def danby_coe_vf(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol): + """ + Vectorized danby_coe + """ + + return danby_coe_hf(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol) + + +@hjit("Tuple([V,V])(f,V,V,f,i8,f)") +def danby_rv_hf(k, r0, v0, tof, numiter, rtol): """Kepler solver for both elliptic and parabolic orbits based on Danby's algorithm. @@ -81,9 +108,9 @@ def danby(k, r0, v0, tof, numiter=20, rtol=1e-8): ---------- k : float Standard gravitational parameter of the attractor. - r0 : numpy.ndarray + r0 : tuple[float,float,float] Position vector. - v0 : numpy.ndarray + v0 : tuple[float,float,float] Velocity vector. tof : float Time of flight. @@ -94,9 +121,9 @@ def danby(k, r0, v0, tof, numiter=20, rtol=1e-8): Returns ------- - rr : numpy.ndarray + rr : tuple[float,float,float] Final position vector. - vv : numpy.ndarray + vv : tuple[float,float,float] Final velocity vector. Notes @@ -105,9 +132,26 @@ def danby(k, r0, v0, tof, numiter=20, rtol=1e-8): Equation* with DOI: https://doi.org/10.1007/BF01686811 """ # Solve first for eccentricity and mean anomaly - p, ecc, inc, raan, argp, nu = rv2coe_hf( - k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL - ) - nu = danby_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol) + p, ecc, inc, raan, argp, nu = rv2coe_hf(k, r0, v0, RV2COE_TOL) + nu = danby_coe_hf(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol) + + return coe2rv_hf(k, p, ecc, inc, raan, argp, nu) + - return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) +@gjit( + "void(f,f[:],f[:],f,i8,f,f[:],f[:])", + "(),(n),(n),(),(),()->(n),(n)", +) +def danby_rv_gf(k, r0, v0, tof, numiter, rtol, rr, vv): + """ + Vectorized danby_rv + """ + + (rr[0], rr[1], rr[2]), (vv[0], vv[1], vv[2]) = danby_rv_hf( + k, + array_to_V_hf(r0), + array_to_V_hf(v0), + tof, + numiter, + rtol, + ) diff --git a/src/hapsira/twobody/propagation/danby.py b/src/hapsira/twobody/propagation/danby.py index e8c4c3d43..1d4e75458 100644 --- a/src/hapsira/twobody/propagation/danby.py +++ b/src/hapsira/twobody/propagation/danby.py @@ -2,7 +2,7 @@ from astropy import units as u -from hapsira.core.propagation import danby_coe as danby_fast +from hapsira.core.propagation.danby import danby_coe_vf, DANBY_NUMITER, DANBY_RTOL from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import ClassicalState @@ -27,10 +27,12 @@ def propagate(self, state, tof): state = state.to_classical() nu = ( - danby_fast( + danby_coe_vf( state.attractor.k.to_value(u.km**3 / u.s**2), *state.to_value(), tof.to_value(u.s), + DANBY_NUMITER, + DANBY_RTOL, ) << u.rad ) diff --git a/tests/tests_core/test_core_propagation.py b/tests/tests_core/test_core_propagation.py index 3c9f7e100..a724278a8 100644 --- a/tests/tests_core/test_core_propagation.py +++ b/tests/tests_core/test_core_propagation.py @@ -3,12 +3,12 @@ import pytest from hapsira.core.propagation import ( - danby_coe, gooding_coe, markley_coe, mikkola_coe, pimienta_coe, ) +from hapsira.core.propagation.danby import danby_coe_vf, DANBY_NUMITER, DANBY_RTOL from hapsira.core.propagation.farnocchia import farnocchia_coe from hapsira.examples import iss @@ -16,7 +16,7 @@ @pytest.mark.parametrize( "propagator_coe", [ - danby_coe, + lambda *args: danby_coe_vf(*args, DANBY_NUMITER, DANBY_RTOL), markley_coe, pimienta_coe, mikkola_coe, From e9ec7fff8dc35822da1276695663fccbd53b7330 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 6 Jan 2024 21:21:02 +0100 Subject: [PATCH 064/346] fix: jit caching in parallel segfaults --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 40619fa4c..0fa102438 100644 --- a/Makefile +++ b/Makefile @@ -49,6 +49,6 @@ upload: done test: - DISPLAY= tox -e style,tests-fast,tests-slow,docs + DISPLAY= HAPSIRA_CACHE=0 tox -e style,tests-fast,tests-slow,docs .PHONY: docs docker image release upload From 30fc77711a69d471807283853a8a13092e767e3d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 6 Jan 2024 21:22:09 +0100 Subject: [PATCH 065/346] jit farnocchia --- src/hapsira/core/propagation/__init__.py | 12 +- src/hapsira/core/propagation/farnocchia.py | 205 +++++++++++------- src/hapsira/twobody/elements.py | 11 +- src/hapsira/twobody/propagation/farnocchia.py | 16 +- tests/tests_core/test_core_propagation.py | 4 +- 5 files changed, 149 insertions(+), 99 deletions(-) diff --git a/src/hapsira/core/propagation/__init__.py b/src/hapsira/core/propagation/__init__.py index 93322cce4..e5023aceb 100644 --- a/src/hapsira/core/propagation/__init__.py +++ b/src/hapsira/core/propagation/__init__.py @@ -4,10 +4,10 @@ from hapsira.core.propagation.cowell import cowell # from hapsira.core.propagation.danby import danby, danby_coe -from hapsira.core.propagation.farnocchia import ( - farnocchia_coe, - farnocchia_rv as farnocchia, -) +# from hapsira.core.propagation.farnocchia import ( +# farnocchia_coe, +# farnocchia_rv as farnocchia, +# ) from hapsira.core.propagation.gooding import gooding, gooding_coe from hapsira.core.propagation.markley import markley, markley_coe from hapsira.core.propagation.mikkola import mikkola, mikkola_coe @@ -18,8 +18,8 @@ __all__ = [ "cowell", "func_twobody", - "farnocchia_coe", - "farnocchia", + # "farnocchia_coe", + # "farnocchia", "vallado", "mikkola_coe", "mikkola", diff --git a/src/hapsira/core/propagation/farnocchia.py b/src/hapsira/core/propagation/farnocchia.py index 9faf5268b..e135c13c2 100644 --- a/src/hapsira/core/propagation/farnocchia.py +++ b/src/hapsira/core/propagation/farnocchia.py @@ -1,7 +1,6 @@ -from numba import njit as jit -import numpy as np +from math import acos, acosh, cos, cosh, nan, pi, sqrt -from hapsira.core.angles import ( +from ..angles import ( D_to_M_hf, D_to_nu_hf, E_to_M_hf, @@ -15,25 +14,34 @@ nu_to_E_hf, nu_to_F_hf, ) -from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL -from ..jit import array_to_V_hf +from ..elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf, hjit, vjit, gjit -@jit -def _kepler_equation_near_parabolic(D, M, ecc): - return D_to_M_near_parabolic(D, ecc) - M +__all__ = [ + "delta_t_from_nu_hf", + "delta_t_from_nu_vf", + "farnocchia_coe_hf", + "farnocchia_coe_vf", + "farnocchia_rv_hf", + "farnocchia_rv_gf", + "FARNOCCHIA_K", + "FARNOCCHIA_Q", + "FARNOCCHIA_DELTA", +] -@jit -def _kepler_equation_prime_near_parabolic(D, M, ecc): - x = (ecc - 1.0) / (ecc + 1.0) * (D**2) - assert abs(x) < 1 - S = dS_x_alt(ecc, x) - return np.sqrt(2.0 / (1.0 + ecc)) + np.sqrt(2.0 / (1.0 + ecc) ** 3) * (D**2) * S +FARNOCCHIA_K = 1.0 +FARNOCCHIA_Q = 1.0 +FARNOCCHIA_DELTA = 1e-2 +_ATOL = 1e-12 +_TOL = 1.48e-08 +_MAXITER = 50 -@jit -def S_x(ecc, x, atol=1e-12): + +@hjit("f(f,f,f)") +def _S_x_hf(ecc, x, atol): assert abs(x) < 1 S = 0 k = 0 @@ -45,8 +53,8 @@ def S_x(ecc, x, atol=1e-12): return S -@jit -def dS_x_alt(ecc, x, atol=1e-12): +@hjit("f(f,f,f)") +def _dS_x_alt_hf(ecc, x, atol): # Notice that this is not exactly # the partial derivative of S with respect to D, # but the result of arranging the terms @@ -62,8 +70,8 @@ def dS_x_alt(ecc, x, atol=1e-12): return S -@jit -def d2S_x_alt(ecc, x, atol=1e-12): +@hjit("f(f,f,f)") +def _d2S_x_alt_hf(ecc, x, atol): # Notice that this is not exactly # the second partial derivative of S with respect to D, # but the result of arranging the terms @@ -80,18 +88,29 @@ def d2S_x_alt(ecc, x, atol=1e-12): return S -@jit -def D_to_M_near_parabolic(D, ecc): +@hjit("f(f,f,f)") +def _kepler_equation_prime_near_parabolic_hf(D, M, ecc): x = (ecc - 1.0) / (ecc + 1.0) * (D**2) assert abs(x) < 1 - S = S_x(ecc, x) - return ( - np.sqrt(2.0 / (1.0 + ecc)) * D + np.sqrt(2.0 / (1.0 + ecc) ** 3) * (D**3) * S - ) + S = _dS_x_alt_hf(ecc, x, _ATOL) + return sqrt(2.0 / (1.0 + ecc)) + sqrt(2.0 / (1.0 + ecc) ** 3) * (D**2) * S + + +@hjit("f(f,f)") +def _D_to_M_near_parabolic_hf(D, ecc): + x = (ecc - 1.0) / (ecc + 1.0) * (D**2) + assert abs(x) < 1 + S = _S_x_hf(ecc, x, _ATOL) + return sqrt(2.0 / (1.0 + ecc)) * D + sqrt(2.0 / (1.0 + ecc) ** 3) * (D**3) * S -@jit -def M_to_D_near_parabolic(M, ecc, tol=1.48e-08, maxiter=50): +@hjit("f(f,f,f)") +def _kepler_equation_near_parabolic_hf(D, M, ecc): + return _D_to_M_near_parabolic_hf(D, ecc) - M + + +@hjit("f(f,f,f,i8)") +def _M_to_D_near_parabolic_hf(M, ecc, tol, maxiter): """Parabolic eccentric anomaly from mean anomaly, near parabolic case. Parameters @@ -114,8 +133,8 @@ def M_to_D_near_parabolic(M, ecc, tol=1.48e-08, maxiter=50): D0 = M_to_D_hf(M) for _ in range(maxiter): - fval = _kepler_equation_near_parabolic(D0, M, ecc) - fder = _kepler_equation_prime_near_parabolic(D0, M, ecc) + fval = _kepler_equation_near_parabolic_hf(D0, M, ecc) + fder = _kepler_equation_prime_near_parabolic_hf(D0, M, ecc) newton_step = fval / fder D = D0 - newton_step @@ -124,11 +143,11 @@ def M_to_D_near_parabolic(M, ecc, tol=1.48e-08, maxiter=50): D0 = D - return np.nan + return nan -@jit -def delta_t_from_nu(nu, ecc, k=1.0, q=1.0, delta=1e-2): +@hjit("f(f,f,f,f,f)") +def delta_t_from_nu_hf(nu, ecc, k, q, delta): """Time elapsed since periapsis for given true anomaly. Parameters @@ -150,18 +169,18 @@ def delta_t_from_nu(nu, ecc, k=1.0, q=1.0, delta=1e-2): Time elapsed since periapsis. """ - assert -np.pi <= nu < np.pi + assert -pi <= nu < pi if ecc < 1 - delta: # Strong elliptic E = nu_to_E_hf(nu, ecc) # (-pi, pi] M = E_to_M_hf(E, ecc) # (-pi, pi] - n = np.sqrt(k * (1 - ecc) ** 3 / q**3) + n = sqrt(k * (1 - ecc) ** 3 / q**3) elif 1 - delta <= ecc < 1: E = nu_to_E_hf(nu, ecc) # (-pi, pi] - if delta <= 1 - ecc * np.cos(E): + if delta <= 1 - ecc * cos(E): # Strong elliptic M = E_to_M_hf(E, ecc) # (-pi, pi] - n = np.sqrt(k * (1 - ecc) ** 3 / q**3) + n = sqrt(k * (1 - ecc) ** 3 / q**3) else: # Near parabolic D = nu_to_D_hf(nu) # (-∞, ∞) @@ -169,43 +188,52 @@ def delta_t_from_nu(nu, ecc, k=1.0, q=1.0, delta=1e-2): # because the near parabolic region shrinks in its vicinity, # otherwise the eccentricity is very close to 1 # and we are really far away - M = D_to_M_near_parabolic(D, ecc) - n = np.sqrt(k / (2 * q**3)) + M = _D_to_M_near_parabolic_hf(D, ecc) + n = sqrt(k / (2 * q**3)) elif ecc == 1: # Parabolic D = nu_to_D_hf(nu) # (-∞, ∞) M = D_to_M_hf(D) # (-∞, ∞) - n = np.sqrt(k / (2 * q**3)) - elif 1 + ecc * np.cos(nu) < 0: + n = sqrt(k / (2 * q**3)) + elif 1 + ecc * cos(nu) < 0: # Unfeasible region - return np.nan + return nan elif 1 < ecc <= 1 + delta: # NOTE: Do we need to wrap nu here? # For hyperbolic orbits, it should anyway be in # (-arccos(-1 / ecc), +arccos(-1 / ecc)) F = nu_to_F_hf(nu, ecc) # (-∞, ∞) - if delta <= ecc * np.cosh(F) - 1: + if delta <= ecc * cosh(F) - 1: # Strong hyperbolic M = F_to_M_hf(F, ecc) # (-∞, ∞) - n = np.sqrt(k * (ecc - 1) ** 3 / q**3) + n = sqrt(k * (ecc - 1) ** 3 / q**3) else: # Near parabolic D = nu_to_D_hf(nu) # (-∞, ∞) - M = D_to_M_near_parabolic(D, ecc) # (-∞, ∞) - n = np.sqrt(k / (2 * q**3)) + M = _D_to_M_near_parabolic_hf(D, ecc) # (-∞, ∞) + n = sqrt(k / (2 * q**3)) elif 1 + delta < ecc: # Strong hyperbolic F = nu_to_F_hf(nu, ecc) # (-∞, ∞) M = F_to_M_hf(F, ecc) # (-∞, ∞) - n = np.sqrt(k * (ecc - 1) ** 3 / q**3) + n = sqrt(k * (ecc - 1) ** 3 / q**3) else: raise RuntimeError return M / n -@jit -def nu_from_delta_t(delta_t, ecc, k=1.0, q=1.0, delta=1e-2): +@vjit("f(f,f,f,f,f)") +def delta_t_from_nu_vf(nu, ecc, k, q, delta): + """ + Vectorized delta_t_from_nu + """ + + return delta_t_from_nu_hf(nu, ecc, k, q, delta) + + +@hjit("f(f,f,f,f,f)") +def _nu_from_delta_t_hf(delta_t, ecc, k, q, delta): """True anomaly for given elapsed time since periapsis. Parameters @@ -229,42 +257,42 @@ def nu_from_delta_t(delta_t, ecc, k=1.0, q=1.0, delta=1e-2): """ if ecc < 1 - delta: # Strong elliptic - n = np.sqrt(k * (1 - ecc) ** 3 / q**3) + n = sqrt(k * (1 - ecc) ** 3 / q**3) M = n * delta_t # This might represent several revolutions, # so we wrap the true anomaly - E = M_to_E_hf((M + np.pi) % (2 * np.pi) - np.pi, ecc) + E = M_to_E_hf((M + pi) % (2 * pi) - pi, ecc) nu = E_to_nu_hf(E, ecc) elif 1 - delta <= ecc < 1: - E_delta = np.arccos((1 - delta) / ecc) + E_delta = acos((1 - delta) / ecc) # We compute M assuming we are in the strong elliptic case # and verify later - n = np.sqrt(k * (1 - ecc) ** 3 / q**3) + n = sqrt(k * (1 - ecc) ** 3 / q**3) M = n * delta_t # We check against abs(M) because E_delta could also be negative if E_to_M_hf(E_delta, ecc) <= abs(M): # Strong elliptic, proceed # This might represent several revolutions, # so we wrap the true anomaly - E = M_to_E_hf((M + np.pi) % (2 * np.pi) - np.pi, ecc) + E = M_to_E_hf((M + pi) % (2 * pi) - pi, ecc) nu = E_to_nu_hf(E, ecc) else: # Near parabolic, recompute M - n = np.sqrt(k / (2 * q**3)) + n = sqrt(k / (2 * q**3)) M = n * delta_t - D = M_to_D_near_parabolic(M, ecc) + D = _M_to_D_near_parabolic_hf(M, ecc, _TOL, _MAXITER) nu = D_to_nu_hf(D) elif ecc == 1: # Parabolic - n = np.sqrt(k / (2 * q**3)) + n = sqrt(k / (2 * q**3)) M = n * delta_t D = M_to_D_hf(M) nu = D_to_nu_hf(D) elif 1 < ecc <= 1 + delta: - F_delta = np.arccosh((1 + delta) / ecc) + F_delta = acosh((1 + delta) / ecc) # We compute M assuming we are in the strong hyperbolic case # and verify later - n = np.sqrt(k * (ecc - 1) ** 3 / q**3) + n = sqrt(k * (ecc - 1) ** 3 / q**3) M = n * delta_t # We check against abs(M) because F_delta could also be negative if F_to_M_hf(F_delta, ecc) <= abs(M): @@ -273,14 +301,14 @@ def nu_from_delta_t(delta_t, ecc, k=1.0, q=1.0, delta=1e-2): nu = F_to_nu_hf(F, ecc) else: # Near parabolic, recompute M - n = np.sqrt(k / (2 * q**3)) + n = sqrt(k / (2 * q**3)) M = n * delta_t - D = M_to_D_near_parabolic(M, ecc) + D = _M_to_D_near_parabolic_hf(M, ecc, _TOL, _MAXITER) nu = D_to_nu_hf(D) # elif 1 + delta < ecc: else: # Strong hyperbolic - n = np.sqrt(k * (ecc - 1) ** 3 / q**3) + n = sqrt(k * (ecc - 1) ** 3 / q**3) M = n * delta_t F = M_to_F_hf(M, ecc) nu = F_to_nu_hf(F, ecc) @@ -288,18 +316,31 @@ def nu_from_delta_t(delta_t, ecc, k=1.0, q=1.0, delta=1e-2): return nu -@jit -def farnocchia_coe(k, p, ecc, inc, raan, argp, nu, tof): - q = p / (1 + ecc) +@hjit("f(f,f,f,f,f,f,f,f)") +def farnocchia_coe_hf(k, p, ecc, inc, raan, argp, nu, tof): + """ + Scalar farnocchia_coe + """ + + q = p / (1.0 + ecc) - delta_t0 = delta_t_from_nu(nu, ecc, k, q) + delta_t0 = delta_t_from_nu_hf(nu, ecc, k, q, FARNOCCHIA_DELTA) delta_t = delta_t0 + tof - return nu_from_delta_t(delta_t, ecc, k, q) + return _nu_from_delta_t_hf(delta_t, ecc, k, q, FARNOCCHIA_DELTA) + + +@vjit("f(f,f,f,f,f,f,f,f)") +def farnocchia_coe_vf(k, p, ecc, inc, raan, argp, nu, tof): + """ + Vectorized farnocchia_coe + """ + return farnocchia_coe_hf(k, p, ecc, inc, raan, argp, nu, tof) -@jit -def farnocchia_rv(k, r0, v0, tof): + +@hjit("Tuple([V,V])(f,V,V,f)") +def farnocchia_rv_hf(k, r0, v0, tof): r"""Propagates orbit using mean motion. This algorithm depends on the geometric shape of the orbit. @@ -315,9 +356,9 @@ def farnocchia_rv(k, r0, v0, tof): ---------- k : float Standar Gravitational parameter - r0 : numpy.ndarray + r0 : tuple[float,float,float] Initial position vector wrt attractor center. - v0 : numpy.ndarray + v0 : tuple[float,float,float] Initial velocity vector. tof : float Time of flight (s). @@ -330,9 +371,21 @@ def farnocchia_rv(k, r0, v0, tof): """ # get the initial true anomaly and orbit parameters that are constant over time - p, ecc, inc, raan, argp, nu0 = rv2coe_hf( - k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL - ) - nu = farnocchia_coe(k, p, ecc, inc, raan, argp, nu0, tof) + p, ecc, inc, raan, argp, nu0 = rv2coe_hf(k, r0, v0, RV2COE_TOL) + nu = farnocchia_coe_hf(k, p, ecc, inc, raan, argp, nu0, tof) - return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) + return coe2rv_hf(k, p, ecc, inc, raan, argp, nu) + + +@gjit( + "void(f,f[:],f[:],f,f[:],f[:])", + "(),(n),(n),()->(n),(n)", +) +def farnocchia_rv_gf(k, r0, v0, tof, rr, vv): + """ + Vectorized farnocchia_rv + """ + + (rr[0], rr[1], rr[2]), (vv[0], vv[1], vv[2]) = farnocchia_rv_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), tof + ) diff --git a/src/hapsira/twobody/elements.py b/src/hapsira/twobody/elements.py index c75b6c8c9..b6d05bd7a 100644 --- a/src/hapsira/twobody/elements.py +++ b/src/hapsira/twobody/elements.py @@ -6,9 +6,7 @@ coe2rv_gf, eccentricity_vector_gf, ) -from hapsira.core.propagation.farnocchia import ( - delta_t_from_nu as delta_t_from_nu_fast, -) +from hapsira.core.propagation.farnocchia import delta_t_from_nu_vf, FARNOCCHIA_DELTA u_kms = u.km / u.s u_km3s2 = u.km**3 / u.s**2 @@ -42,9 +40,9 @@ def energy(k, r, v): @u.quantity_input(k=u_km3s2, r=u.km, v=u_kms) def eccentricity_vector(k, r, v): """Eccentricity vector.""" - e = eccentricity_vector_gf( + e = eccentricity_vector_gf( # pylint: disable=E1120 k.to_value(u_km3s2), r.to_value(u.km), v.to_value(u_kms) - ) # pylint: disable=E1120 + ) return e << u.one @@ -53,11 +51,12 @@ def t_p(nu, ecc, k, r_p): """Elapsed time since latest perifocal passage.""" # TODO: Make this a propagator method t_p = ( - delta_t_from_nu_fast( + delta_t_from_nu_vf( nu.to_value(u.rad), ecc.value, k.to_value(u_km3s2), r_p.to_value(u.km), + FARNOCCHIA_DELTA, ) * u.s ) diff --git a/src/hapsira/twobody/propagation/farnocchia.py b/src/hapsira/twobody/propagation/farnocchia.py index 4b26f57d6..d499a41fd 100644 --- a/src/hapsira/twobody/propagation/farnocchia.py +++ b/src/hapsira/twobody/propagation/farnocchia.py @@ -1,11 +1,10 @@ import sys from astropy import units as u -import numpy as np from hapsira.core.propagation.farnocchia import ( - farnocchia_coe as farnocchia_coe_fast, - farnocchia_rv as farnocchia_rv_fast, + farnocchia_coe_vf, + farnocchia_rv_gf, ) from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import ClassicalState @@ -34,7 +33,7 @@ def propagate(self, state, tof): state = state.to_classical() nu = ( - farnocchia_coe_fast( + farnocchia_coe_vf( state.attractor.k.to_value(u.km**3 / u.s**2), *state.to_value(), tof.to_value(u.s), @@ -54,10 +53,9 @@ def propagate_many(self, state, tofs): # TODO: This should probably return a ClassicalStateArray instead, # see discussion at https://github.com/hapsira/hapsira/pull/1492 - results = np.array( - [farnocchia_rv_fast(k, *rv0, tof) for tof in tofs.to_value(u.s)] - ) + rr, vv = farnocchia_rv_gf(k, *rv0, tofs.to_value(u.s)) # pylint: disable=E0633 + return ( - results[:, 0] << u.km, - results[:, 1] << (u.km / u.s), + rr << u.km, + vv << (u.km / u.s), ) diff --git a/tests/tests_core/test_core_propagation.py b/tests/tests_core/test_core_propagation.py index a724278a8..b2f77df36 100644 --- a/tests/tests_core/test_core_propagation.py +++ b/tests/tests_core/test_core_propagation.py @@ -9,7 +9,7 @@ pimienta_coe, ) from hapsira.core.propagation.danby import danby_coe_vf, DANBY_NUMITER, DANBY_RTOL -from hapsira.core.propagation.farnocchia import farnocchia_coe +from hapsira.core.propagation.farnocchia import farnocchia_coe_vf from hapsira.examples import iss @@ -20,7 +20,7 @@ markley_coe, pimienta_coe, mikkola_coe, - farnocchia_coe, + farnocchia_coe_vf, gooding_coe, ], ) From 27f3ddd247a27e69b9bae2c59c04783189c73533 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 7 Jan 2024 11:42:28 +0100 Subject: [PATCH 066/346] jit gooding --- src/hapsira/core/propagation/__init__.py | 6 +- src/hapsira/core/propagation/gooding.py | 82 +++++++++++++++------- src/hapsira/twobody/propagation/gooding.py | 10 ++- tests/tests_core/test_core_propagation.py | 8 ++- 4 files changed, 75 insertions(+), 31 deletions(-) diff --git a/src/hapsira/core/propagation/__init__.py b/src/hapsira/core/propagation/__init__.py index e5023aceb..b2982287b 100644 --- a/src/hapsira/core/propagation/__init__.py +++ b/src/hapsira/core/propagation/__init__.py @@ -8,7 +8,7 @@ # farnocchia_coe, # farnocchia_rv as farnocchia, # ) -from hapsira.core.propagation.gooding import gooding, gooding_coe +# from hapsira.core.propagation.gooding import gooding, gooding_coe from hapsira.core.propagation.markley import markley, markley_coe from hapsira.core.propagation.mikkola import mikkola, mikkola_coe from hapsira.core.propagation.pimienta import pimienta, pimienta_coe @@ -27,8 +27,8 @@ "markley", "pimienta_coe", "pimienta", - "gooding_coe", - "gooding", + # "gooding_coe", + # "gooding", # "danby_coe", # "danby", "recseries_coe", diff --git a/src/hapsira/core/propagation/gooding.py b/src/hapsira/core/propagation/gooding.py index 51c458f5f..bd290f394 100644 --- a/src/hapsira/core/propagation/gooding.py +++ b/src/hapsira/core/propagation/gooding.py @@ -1,13 +1,29 @@ -from numba import njit as jit -import numpy as np +from math import cos, sin, sqrt -from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf -from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL -from ..jit import array_to_V_hf +from ..angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf +from ..elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf, hjit, vjit, gjit -@jit -def gooding_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter=150, rtol=1e-8): +__all__ = [ + "gooding_coe_hf", + "gooding_coe_vf", + "gooding_rv_hf", + "gooding_rv_gf", + "GOODING_NUMITER", + "GOODING_RTOL", +] + + +GOODING_NUMITER = 150 +GOODING_RTOL = 1e-8 + + +@hjit("f(f,f,f,f,f,f,f,f,i8,f)") +def gooding_coe_hf(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol): + """ + Scalar gooding_coe + """ # TODO: parabolic and hyperbolic not implemented cases if ecc >= 1.0: raise NotImplementedError( @@ -16,18 +32,18 @@ def gooding_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter=150, rtol=1e-8): M0 = E_to_M_hf(nu_to_E_hf(nu, ecc), ecc) semi_axis_a = p / (1 - ecc**2) - n = np.sqrt(k / np.abs(semi_axis_a) ** 3) + n = sqrt(k / abs(semi_axis_a) ** 3) M = M0 + n * tof # Start the computation n = 0 - c = ecc * np.cos(M) - s = ecc * np.sin(M) - psi = s / np.sqrt(1 - 2 * c + ecc**2) + c = ecc * cos(M) + s = ecc * sin(M) + psi = s / sqrt(1 - 2 * c + ecc**2) f = 1.0 while f**2 >= rtol and n <= numiter: - xi = np.cos(psi) - eta = np.sin(psi) + xi = cos(psi) + eta = sin(psi) fd = (1 - c * xi) + s * eta fdd = c * eta + s * xi f = psi - fdd @@ -38,8 +54,17 @@ def gooding_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter=150, rtol=1e-8): return E_to_nu_hf(E, ecc) -@jit -def gooding(k, r0, v0, tof, numiter=150, rtol=1e-8): +@vjit("f(f,f,f,f,f,f,f,f,i8,f)") +def gooding_coe_vf(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol): + """ + Vectorized gooding_coe + """ + + return gooding_coe_hf(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol) + + +@hjit("Tuple([V,V])(f,V,V,f,i8,f)") +def gooding_rv_hf(k, r0, v0, tof, numiter, rtol): """Solves the Elliptic Kepler Equation with a cubic convergence and accuracy better than 10e-12 rad is normally achieved. It is not valid for eccentricities equal or higher than 1.0. @@ -48,9 +73,9 @@ def gooding(k, r0, v0, tof, numiter=150, rtol=1e-8): ---------- k : float Standard gravitational parameter of the attractor. - r0 : numpy.ndarray + r0 : tuple[float,float,float] Position vector. - v0 : numpy.ndarray + v0 : tuple[float,float,float] Velocity vector. tof : float Time of flight. @@ -61,18 +86,27 @@ def gooding(k, r0, v0, tof, numiter=150, rtol=1e-8): Returns ------- - rr : numpy.ndarray + rr : tuple[float,float,float] Final position vector. - vv : numpy.ndarray + vv : tuple[float,float,float] Final velocity vector. Note ---- Original paper for the algorithm: https://doi.org/10.1007/BF01238923 """ # Solve first for eccentricity and mean anomaly - p, ecc, inc, raan, argp, nu = rv2coe_hf( - k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL - ) - nu = gooding_coe(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol) + p, ecc, inc, raan, argp, nu = rv2coe_hf(k, r0, v0, RV2COE_TOL) + nu = gooding_coe_hf(k, p, ecc, inc, raan, argp, nu, tof, numiter, rtol) + + return coe2rv_hf(k, p, ecc, inc, raan, argp, nu) + - return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) +@gjit("void(f,f[:],f[:],f,i8,f,f[:],f[:])", "(),(n),(n),(),(),()->(),()") +def gooding_rv_gf(k, r0, v0, tof, numiter, rtol, rr, vv): + """ + Vectorized gooding_rv + """ + + (rr[0], rr[1], rr[2]), (vv[0], vv[1], vv[2]) = gooding_rv_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), tof, numiter, rtol + ) diff --git a/src/hapsira/twobody/propagation/gooding.py b/src/hapsira/twobody/propagation/gooding.py index 3c0da8fcd..11f9eed75 100644 --- a/src/hapsira/twobody/propagation/gooding.py +++ b/src/hapsira/twobody/propagation/gooding.py @@ -2,7 +2,11 @@ from astropy import units as u -from hapsira.core.propagation import gooding_coe as gooding_fast +from hapsira.core.propagation.gooding import ( + gooding_coe_vf, + GOODING_RTOL, + GOODING_NUMITER, +) from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import ClassicalState @@ -32,10 +36,12 @@ def propagate(self, state, tof): state = state.to_classical() nu = ( - gooding_fast( + gooding_coe_vf( state.attractor.k.to_value(u.km**3 / u.s**2), *state.to_value(), tof.to_value(u.s), + GOODING_NUMITER, + GOODING_RTOL, ) << u.rad ) diff --git a/tests/tests_core/test_core_propagation.py b/tests/tests_core/test_core_propagation.py index b2f77df36..7d23a56cb 100644 --- a/tests/tests_core/test_core_propagation.py +++ b/tests/tests_core/test_core_propagation.py @@ -3,13 +3,17 @@ import pytest from hapsira.core.propagation import ( - gooding_coe, markley_coe, mikkola_coe, pimienta_coe, ) from hapsira.core.propagation.danby import danby_coe_vf, DANBY_NUMITER, DANBY_RTOL from hapsira.core.propagation.farnocchia import farnocchia_coe_vf +from hapsira.core.propagation.gooding import ( + gooding_coe_vf, + GOODING_NUMITER, + GOODING_RTOL, +) from hapsira.examples import iss @@ -21,7 +25,7 @@ pimienta_coe, mikkola_coe, farnocchia_coe_vf, - gooding_coe, + lambda *args: gooding_coe_vf(*args, GOODING_NUMITER, GOODING_RTOL), ], ) def test_propagate_with_coe(propagator_coe): From 8a555f52019ea59b9b40aaedfa25f96c94b76d77 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 8 Jan 2024 08:13:20 +0100 Subject: [PATCH 067/346] div similar to numpy --- src/hapsira/core/math/linalg.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index d7719505e..985fce29e 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -1,4 +1,4 @@ -from math import sqrt +from math import inf, sqrt from ..jit import hjit, vjit @@ -32,6 +32,16 @@ def cross_VV_hf(a, b): ) +@hjit("f(f,f)") +def div_ss_hf(a, b): + """ + Similar to np.divide + """ + if b == 0: + return inf if a >= 0 else -inf + return a / b + + @hjit("V(V,f)") def div_Vs_hf(v, s): return v[0] / s, v[1] / s, v[2] / s From 5efd002a8d347e1b5ee33b513c5ab8b561e587cb Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 8 Jan 2024 08:13:54 +0100 Subject: [PATCH 068/346] div similar to numpy --- src/hapsira/core/iod.py | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/iod.py b/src/hapsira/core/iod.py index 1fef127ce..a17bb2487 100644 --- a/src/hapsira/core/iod.py +++ b/src/hapsira/core/iod.py @@ -4,6 +4,7 @@ from .math.linalg import ( add_VV_hf, cross_VV_hf, + div_ss_hf, div_Vs_hf, matmul_VV_hf, mul_Vs_hf, @@ -24,6 +25,8 @@ "compute_y_vf", "tof_equation_y_hf", "tof_equation_y_vf", + "find_xy_hf", + "find_xy_gf", ] @@ -136,8 +139,10 @@ def tof_equation_y_hf(x, y, T0, ll, M): T_ = (eta**3 * Q + 4 * ll * eta) * 0.5 else: psi = _compute_psi_hf(x, y, ll) - T_ = (((psi + M * pi) / sqrt(abs(1 - x**2))) - x + ll * y) / (1 - x**2) - + T_ = div_ss_hf( + (div_ss_hf((psi + M * pi), sqrt(abs(1 - x**2))) - x + ll * y), + (1 - x**2), + ) return T_ - T0 @@ -245,7 +250,7 @@ def _householder_hf(p0, T0, ll, M, tol, maxiter): @hjit("Tuple([f,f])(f,f,i8,i8,b1,f)") -def _find_xy_hf(ll, T, M, numiter, lowpath, rtol): +def find_xy_hf(ll, T, M, numiter, lowpath, rtol): """Computes all x, y for given number of revolutions.""" # For abs(ll) == 1 the derivative is not continuous assert abs(ll) < 1 @@ -275,6 +280,18 @@ def _find_xy_hf(ll, T, M, numiter, lowpath, rtol): return x, y +@gjit( + "void(f,f,i8,i8,b1,f,f[:],f[:])", + "(),(),(),(),(),()->(),()", +) +def find_xy_gf(ll, T, M, numiter, lowpath, rtol, x, y): + """ + Vectorized find_xy + """ + + x[0], y[0] = find_xy_hf(ll, T, M, numiter, lowpath, rtol) + + @hjit("Tuple([V,V])(f,V,V,f,i8,b1,b1,i8,f)") def vallado_hf(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): r"""Solves the Lambert's problem. @@ -538,7 +555,7 @@ def izzo_hf(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol): T = sqrt(2 * k / s**3) * tof # Find solutions - x, y = _find_xy_hf(ll, T, M, numiter, lowpath, rtol) + x, y = find_xy_hf(ll, T, M, numiter, lowpath, rtol) # Reconstruct gamma = sqrt(k * s / 2) From a7feb50963abd2fb513fd36aaaa58ce03da975a3 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 8 Jan 2024 08:14:11 +0100 Subject: [PATCH 069/346] use new core api --- ...isiting-lamberts-problem-in-python.myst.md | 48 ++++++++++--------- 1 file changed, 26 insertions(+), 22 deletions(-) diff --git a/docs/source/examples/revisiting-lamberts-problem-in-python.myst.md b/docs/source/examples/revisiting-lamberts-problem-in-python.myst.md index d931ea226..b96a1cf6f 100644 --- a/docs/source/examples/revisiting-lamberts-problem-in-python.myst.md +++ b/docs/source/examples/revisiting-lamberts-problem-in-python.myst.md @@ -4,9 +4,9 @@ jupytext: extension: .md format_name: myst format_version: 0.13 - jupytext_version: 1.14.0 + jupytext_version: 1.16.0 kernelspec: - display_name: Python 3 + display_name: Python 3 (ipykernel) language: python name: python3 --- @@ -15,25 +15,25 @@ kernelspec: The Izzo algorithm to solve the Lambert problem is available in hapsira and was implemented from [this paper](https://arxiv.org/abs/1403.2705). -```{code-cell} +```{code-cell} ipython3 from cycler import cycler from matplotlib import pyplot as plt import numpy as np -from hapsira.core import iod +from hapsira.core.iod import compute_y_vf, tof_equation_y_vf, compute_T_min_gf, find_xy_gf from hapsira.iod import izzo ``` ## Part 1: Reproducing the original figure -```{code-cell} +```{code-cell} ipython3 x = np.linspace(-1, 2, num=1000) M_list = 0, 1, 2, 3 ll_list = 1, 0.9, 0.7, 0, -0.7, -0.9, -1 ``` -```{code-cell} +```{code-cell} ipython3 fig, ax = plt.subplots(figsize=(10, 8)) ax.set_prop_cycle( cycler("linestyle", ["-", "--"]) @@ -43,8 +43,8 @@ for M in M_list: for ll in ll_list: T_x0 = np.zeros_like(x) for ii in range(len(x)): - y = iod._compute_y(x[ii], ll) - T_x0[ii] = iod._tof_equation_y(x[ii], y, 0.0, ll, M) + y = compute_y_vf(x[ii], ll) + T_x0[ii] = tof_equation_y_vf(x[ii], y, 0.0, ll, M) if M == 0 and ll == 1: T_x0[x > 0] = np.nan elif M > 0: @@ -86,12 +86,12 @@ ax.set_ylabel("$T$") ## Part 2: Locating $T_{min}$ -```{code-cell} +```{code-cell} ipython3 :tags: [nbsphinx-thumbnail] for M in M_list: for ll in ll_list: - x_T_min, T_min = iod._compute_T_min(ll, M, 10, 1e-8) + x_T_min, T_min = compute_T_min_gf(ll, M, 10, 1e-8) ax.plot(x_T_min, T_min, "kx", mew=2) fig @@ -99,17 +99,21 @@ fig ## Part 3: Try out solution -```{code-cell} +```{code-cell} ipython3 T_ref = 1 ll_ref = 0 -x_ref, _ = iod._find_xy( - ll_ref, T_ref, M=0, numiter=10, lowpath=True, rtol=1e-8 +x_ref, _ = find_xy_gf( + ll_ref, T_ref, + 0, # M + 10, # numiter + True, # lowpath + 1e-8, # rtol ) x_ref ``` -```{code-cell} +```{code-cell} ipython3 ax.plot(x_ref, T_ref, "o", mew=2, mec="red", mfc="none") fig @@ -117,7 +121,7 @@ fig ## Part 4: Run some examples -```{code-cell} +```{code-cell} ipython3 from astropy import units as u from hapsira.bodies import Earth @@ -125,7 +129,7 @@ from hapsira.bodies import Earth ### Single revolution -```{code-cell} +```{code-cell} ipython3 k = Earth.k r0 = [15945.34, 0.0, 0.0] * u.km r = [12214.83399, 10249.46731, 0.0] * u.km @@ -138,7 +142,7 @@ v0, v = izzo.lambert(k, r0, r, tof) v ``` -```{code-cell} +```{code-cell} ipython3 k = Earth.k r0 = [5000.0, 10000.0, 2100.0] * u.km r = [-14600.0, 2500.0, 7000.0] * u.km @@ -153,7 +157,7 @@ v ### Multiple revolutions -```{code-cell} +```{code-cell} ipython3 k = Earth.k r0 = [22592.145603, -1599.915239, -19783.950506] * u.km r = [1922.067697, 4054.157051, -8925.727465] * u.km @@ -169,20 +173,20 @@ expected_va_r = [-2.45759553, 1.16945801, 0.43161258] * u.km / u.s expected_vb_r = [-5.53841370, 0.01822220, 5.49641054] * u.km / u.s ``` -```{code-cell} +```{code-cell} ipython3 v0, v = izzo.lambert(k, r0, r, tof, M=0) v ``` -```{code-cell} +```{code-cell} ipython3 _, v_l = izzo.lambert(k, r0, r, tof, M=1, lowpath=True) _, v_r = izzo.lambert(k, r0, r, tof, M=1, lowpath=False) ``` -```{code-cell} +```{code-cell} ipython3 v_l ``` -```{code-cell} +```{code-cell} ipython3 v_r ``` From 726819c86b2d868d9f8edcde740027b428776bf7 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 8 Jan 2024 15:53:19 +0100 Subject: [PATCH 070/346] jit markley --- src/hapsira/core/propagation/__init__.py | 6 +- src/hapsira/core/propagation/markley.py | 77 +++++++++++++++------- src/hapsira/twobody/propagation/markley.py | 4 +- tests/tests_core/test_core_propagation.py | 4 +- 4 files changed, 60 insertions(+), 31 deletions(-) diff --git a/src/hapsira/core/propagation/__init__.py b/src/hapsira/core/propagation/__init__.py index b2982287b..8a1abb3d7 100644 --- a/src/hapsira/core/propagation/__init__.py +++ b/src/hapsira/core/propagation/__init__.py @@ -9,7 +9,7 @@ # farnocchia_rv as farnocchia, # ) # from hapsira.core.propagation.gooding import gooding, gooding_coe -from hapsira.core.propagation.markley import markley, markley_coe +# from hapsira.core.propagation.markley import markley, markley_coe from hapsira.core.propagation.mikkola import mikkola, mikkola_coe from hapsira.core.propagation.pimienta import pimienta, pimienta_coe from hapsira.core.propagation.recseries import recseries, recseries_coe @@ -23,8 +23,8 @@ "vallado", "mikkola_coe", "mikkola", - "markley_coe", - "markley", + # "markley_coe", + # "markley", "pimienta_coe", "pimienta", # "gooding_coe", diff --git a/src/hapsira/core/propagation/markley.py b/src/hapsira/core/propagation/markley.py index 5320b6a01..743f70d73 100644 --- a/src/hapsira/core/propagation/markley.py +++ b/src/hapsira/core/propagation/markley.py @@ -1,29 +1,40 @@ -from numba import njit as jit -import numpy as np +from math import cos, pi, sin, sqrt -from hapsira.core.angles import ( +from ..angles import ( E_to_M_hf, E_to_nu_hf, kepler_equation_hf, kepler_equation_prime_hf, nu_to_E_hf, ) -from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL -from ..jit import array_to_V_hf +from ..elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf, hjit, vjit, gjit -@jit -def markley_coe(k, p, ecc, inc, raan, argp, nu, tof): +__all__ = [ + "markley_coe_hf", + "markley_coe_vf", + "markley_rv_hf", + "markley_rv_gf", +] + + +@hjit("f(f,f,f,f,f,f,f,f)") +def markley_coe_hf(k, p, ecc, inc, raan, argp, nu, tof): + """ + Scalar markley_coe + """ + M0 = E_to_M_hf(nu_to_E_hf(nu, ecc), ecc) a = p / (1 - ecc**2) - n = np.sqrt(k / a**3) + n = sqrt(k / a**3) M = M0 + n * tof # Range between -pi and pi - M = (M + np.pi) % (2 * np.pi) - np.pi + M = (M + pi) % (2 * pi) - pi # Equation (20) - alpha = (3 * np.pi**2 + 1.6 * (np.pi - np.abs(M)) / (1 + ecc)) / (np.pi**2 - 6) + alpha = (3 * pi**2 + 1.6 * (pi - abs(M)) / (1 + ecc)) / (pi**2 - 6) # Equation (5) d = 3 * (1 - ecc) + alpha * ecc @@ -35,7 +46,7 @@ def markley_coe(k, p, ecc, inc, raan, argp, nu, tof): r = 3 * alpha * d * (d - 1 + ecc) * M + M**3 # Equation (14) - w = (np.abs(r) + np.sqrt(q**3 + r**2)) ** (2 / 3) + w = (abs(r) + sqrt(q**3 + r**2)) ** (2 / 3) # Equation (15) E = (2 * r * w / (w**2 + w * q + q**2) + M) / d @@ -43,8 +54,8 @@ def markley_coe(k, p, ecc, inc, raan, argp, nu, tof): # Equation (26) f0 = kepler_equation_hf(E, M, ecc) f1 = kepler_equation_prime_hf(E, M, ecc) - f2 = ecc * np.sin(E) - f3 = ecc * np.cos(E) + f2 = ecc * sin(E) + f3 = ecc * cos(E) f4 = -f2 # Equation (22) @@ -60,8 +71,17 @@ def markley_coe(k, p, ecc, inc, raan, argp, nu, tof): return nu -@jit -def markley(k, r0, v0, tof): +@vjit("f(f,f,f,f,f,f,f,f)") +def markley_coe_vf(k, p, ecc, inc, raan, argp, nu, tof): + """ + Vectorized markley_coe + """ + + return markley_coe_hf(k, p, ecc, inc, raan, argp, nu, tof) + + +@hjit("Tuple([V,V])(f,V,V,f)") +def markley_rv_hf(k, r0, v0, tof): """Solves the kepler problem by a non-iterative method. Relative error is around 1e-18, only limited by machine double-precision errors. @@ -69,18 +89,18 @@ def markley(k, r0, v0, tof): ---------- k : float Standar Gravitational parameter. - r0 : numpy.ndarray + r0 : tuple[float,float,float] Initial position vector wrt attractor center. - v0 : numpy.ndarray + v0 : tuple[float,float,float] Initial velocity vector. tof : float Time of flight. Returns ------- - rr: numpy.ndarray + rr: tuple[float,float,float] Final position vector. - vv: numpy.ndarray + vv: tuple[float,float,float] Final velocity vector. Notes @@ -89,9 +109,18 @@ def markley(k, r0, v0, tof): """ # Solve first for eccentricity and mean anomaly - p, ecc, inc, raan, argp, nu = rv2coe_hf( - k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL - ) - nu = markley_coe(k, p, ecc, inc, raan, argp, nu, tof) + p, ecc, inc, raan, argp, nu = rv2coe_hf(k, r0, v0, RV2COE_TOL) + nu = markley_coe_hf(k, p, ecc, inc, raan, argp, nu, tof) + + return coe2rv_hf(k, p, ecc, inc, raan, argp, nu) + - return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) +@gjit("void(f,f[:],f[:],f,f[:],f[:])", "(),(n),(n),()->(n),(n)") +def markley_rv_gf(k, r0, v0, tof, rr, vv): + """ + Vectorized markley_rv + """ + + (rr[0], rr[1], rr[2]), (vv[0], vv[1], vv[2]) = markley_rv_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), tof + ) diff --git a/src/hapsira/twobody/propagation/markley.py b/src/hapsira/twobody/propagation/markley.py index ed99efec2..530ac7a4e 100644 --- a/src/hapsira/twobody/propagation/markley.py +++ b/src/hapsira/twobody/propagation/markley.py @@ -2,7 +2,7 @@ from astropy import units as u -from hapsira.core.propagation import markley_coe as markley_fast +from hapsira.core.propagation.markley import markley_coe_vf from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import ClassicalState @@ -28,7 +28,7 @@ def propagate(self, state, tof): state = state.to_classical() nu = ( - markley_fast( + markley_coe_vf( state.attractor.k.to_value(u.km**3 / u.s**2), *state.to_value(), tof.to_value(u.s), diff --git a/tests/tests_core/test_core_propagation.py b/tests/tests_core/test_core_propagation.py index 7d23a56cb..69e0c571d 100644 --- a/tests/tests_core/test_core_propagation.py +++ b/tests/tests_core/test_core_propagation.py @@ -3,7 +3,6 @@ import pytest from hapsira.core.propagation import ( - markley_coe, mikkola_coe, pimienta_coe, ) @@ -14,6 +13,7 @@ GOODING_NUMITER, GOODING_RTOL, ) +from hapsira.core.propagation.markley import markley_coe_vf from hapsira.examples import iss @@ -21,7 +21,7 @@ "propagator_coe", [ lambda *args: danby_coe_vf(*args, DANBY_NUMITER, DANBY_RTOL), - markley_coe, + markley_coe_vf, pimienta_coe, mikkola_coe, farnocchia_coe_vf, From 4285899bb8df5051db622e23bea57473d0e8d3ac Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 8 Jan 2024 16:12:35 +0100 Subject: [PATCH 071/346] jit mikkola --- src/hapsira/core/propagation/__init__.py | 6 +- src/hapsira/core/propagation/mikkola.py | 89 ++++++++++++++-------- src/hapsira/twobody/propagation/mikkola.py | 4 +- tests/tests_core/test_core_propagation.py | 4 +- 4 files changed, 66 insertions(+), 37 deletions(-) diff --git a/src/hapsira/core/propagation/__init__.py b/src/hapsira/core/propagation/__init__.py index 8a1abb3d7..31f016a39 100644 --- a/src/hapsira/core/propagation/__init__.py +++ b/src/hapsira/core/propagation/__init__.py @@ -10,7 +10,7 @@ # ) # from hapsira.core.propagation.gooding import gooding, gooding_coe # from hapsira.core.propagation.markley import markley, markley_coe -from hapsira.core.propagation.mikkola import mikkola, mikkola_coe +# from hapsira.core.propagation.mikkola import mikkola, mikkola_coe from hapsira.core.propagation.pimienta import pimienta, pimienta_coe from hapsira.core.propagation.recseries import recseries, recseries_coe from hapsira.core.propagation.vallado import vallado @@ -21,8 +21,8 @@ # "farnocchia_coe", # "farnocchia", "vallado", - "mikkola_coe", - "mikkola", + # "mikkola_coe", + # "mikkola", # "markley_coe", # "markley", "pimienta_coe", diff --git a/src/hapsira/core/propagation/mikkola.py b/src/hapsira/core/propagation/mikkola.py index 538e9d8d8..23a5cc6c5 100644 --- a/src/hapsira/core/propagation/mikkola.py +++ b/src/hapsira/core/propagation/mikkola.py @@ -1,7 +1,6 @@ -from numba import njit as jit -import numpy as np +from math import cos, cosh, log, sin, sinh, sqrt -from hapsira.core.angles import ( +from ..angles import ( D_to_nu_hf, E_to_M_hf, E_to_nu_hf, @@ -10,14 +9,26 @@ nu_to_E_hf, nu_to_F_hf, ) -from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL -from ..jit import array_to_V_hf +from ..elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf, hjit, vjit, gjit -@jit -def mikkola_coe(k, p, ecc, inc, raan, argp, nu, tof): +__all__ = [ + "mikkola_coe_hf", + "mikkola_coe_vf", + "mikkola_rv_hf", + "mikkola_rv_gf", +] + + +@hjit("f(f,f,f,f,f,f,f,f)") +def mikkola_coe_hf(k, p, ecc, inc, raan, argp, nu, tof): + """ + Scalar mikkola_coe + """ + a = p / (1 - ecc**2) - n = np.sqrt(k / np.abs(a) ** 3) + n = sqrt(k / abs(a) ** 3) # Solve for specific geometrical case if ecc < 1.0: @@ -33,9 +44,9 @@ def mikkola_coe(k, p, ecc, inc, raan, argp, nu, tof): # Equation (9b) if beta >= 0: - z = (beta + np.sqrt(beta**2 + alpha**3)) ** (1 / 3) + z = (beta + sqrt(beta**2 + alpha**3)) ** (1 / 3) else: - z = (beta - np.sqrt(beta**2 + alpha**3)) ** (1 / 3) + z = (beta - sqrt(beta**2 + alpha**3)) ** (1 / 3) s = z - alpha / z @@ -50,18 +61,18 @@ def mikkola_coe(k, p, ecc, inc, raan, argp, nu, tof): # Solving for the true anomaly if ecc < 1.0: E = M + ecc * (3 * s - 4 * s**3) - f = E - ecc * np.sin(E) - M - f1 = 1.0 - ecc * np.cos(E) - f2 = ecc * np.sin(E) - f3 = ecc * np.cos(E) + f = E - ecc * sin(E) - M + f1 = 1.0 - ecc * cos(E) + f2 = ecc * sin(E) + f3 = ecc * cos(E) f4 = -f2 f5 = -f3 else: - E = 3 * np.log(s + np.sqrt(1 + s**2)) - f = -E + ecc * np.sinh(E) - M - f1 = -1.0 + ecc * np.cosh(E) - f2 = ecc * np.sinh(E) - f3 = ecc * np.cosh(E) + E = 3 * log(s + sqrt(1 + s**2)) + f = -E + ecc * sinh(E) - M + f1 = -1.0 + ecc * cosh(E) + f2 = ecc * sinh(E) + f3 = ecc * cosh(E) f4 = f2 f5 = f3 @@ -95,17 +106,26 @@ def mikkola_coe(k, p, ecc, inc, raan, argp, nu, tof): return nu -@jit -def mikkola(k, r0, v0, tof, rtol=None): +@vjit("f(f,f,f,f,f,f,f,f)") +def mikkola_coe_vf(k, p, ecc, inc, raan, argp, nu, tof): + """ + Vectorized mikkola_coe + """ + + return mikkola_coe_hf(k, p, ecc, inc, raan, argp, nu, tof) + + +@hjit("Tuple([V,V])(f,V,V,f)") +def mikkola_rv_hf(k, r0, v0, tof): # rtol=None """Raw algorithm for Mikkola's Kepler solver. Parameters ---------- k : float Standard gravitational parameter of the attractor. - r0 : numpy.ndarray + r0 : tuple[float,float,float] Position vector. - v0 : numpy.ndarray + v0 : tuple[float,float,float] Velocity vector. tof : float Time of flight. @@ -114,18 +134,27 @@ def mikkola(k, r0, v0, tof, rtol=None): Returns ------- - rr : numpy.ndarray + rr : tuple[float,float,float] Final velocity vector. - vv : numpy.ndarray + vv : tuple[float,float,float] Final velocity vector. Note ---- Original paper: https://doi.org/10.1007/BF01235850 """ # Solving for the classical elements - p, ecc, inc, raan, argp, nu = rv2coe_hf( - k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL - ) - nu = mikkola_coe(k, p, ecc, inc, raan, argp, nu, tof) + p, ecc, inc, raan, argp, nu = rv2coe_hf(k, r0, v0, RV2COE_TOL) + nu = mikkola_coe_hf(k, p, ecc, inc, raan, argp, nu, tof) + + return coe2rv_hf(k, p, ecc, inc, raan, argp, nu) + - return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) +@gjit("void(f,f[:],f[:],f,f[:],f[:])", "(),(n),(n),()->(n),(n)") +def mikkola_rv_gf(k, r0, v0, tof, rr, vv): # rtol=None + """ + Vectorized mikkola_rv + """ + + (rr[0], rr[1], rr[2]), (vv[0], vv[1], vv[2]) = mikkola_rv_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), tof + ) diff --git a/src/hapsira/twobody/propagation/mikkola.py b/src/hapsira/twobody/propagation/mikkola.py index 1f5233af7..a0f999fb8 100644 --- a/src/hapsira/twobody/propagation/mikkola.py +++ b/src/hapsira/twobody/propagation/mikkola.py @@ -2,7 +2,7 @@ from astropy import units as u -from hapsira.core.propagation import mikkola_coe as mikkola_fast +from hapsira.core.propagation.mikkola import mikkola_coe_vf from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import ClassicalState @@ -27,7 +27,7 @@ def propagate(self, state, tof): state = state.to_classical() nu = ( - mikkola_fast( + mikkola_coe_vf( state.attractor.k.to_value(u.km**3 / u.s**2), *state.to_value(), tof.to_value(u.s), diff --git a/tests/tests_core/test_core_propagation.py b/tests/tests_core/test_core_propagation.py index 69e0c571d..114c6c27d 100644 --- a/tests/tests_core/test_core_propagation.py +++ b/tests/tests_core/test_core_propagation.py @@ -3,7 +3,6 @@ import pytest from hapsira.core.propagation import ( - mikkola_coe, pimienta_coe, ) from hapsira.core.propagation.danby import danby_coe_vf, DANBY_NUMITER, DANBY_RTOL @@ -14,6 +13,7 @@ GOODING_RTOL, ) from hapsira.core.propagation.markley import markley_coe_vf +from hapsira.core.propagation.mikkola import mikkola_coe_vf from hapsira.examples import iss @@ -23,7 +23,7 @@ lambda *args: danby_coe_vf(*args, DANBY_NUMITER, DANBY_RTOL), markley_coe_vf, pimienta_coe, - mikkola_coe, + mikkola_coe_vf, farnocchia_coe_vf, lambda *args: gooding_coe_vf(*args, GOODING_NUMITER, GOODING_RTOL), ], From fae6f4de0d4c1f377975e69b5e07ae6166a5f095 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 8 Jan 2024 16:14:14 +0100 Subject: [PATCH 072/346] cleanup --- src/hapsira/core/propagation/mikkola.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/propagation/mikkola.py b/src/hapsira/core/propagation/mikkola.py index 23a5cc6c5..5e4a87605 100644 --- a/src/hapsira/core/propagation/mikkola.py +++ b/src/hapsira/core/propagation/mikkola.py @@ -116,7 +116,7 @@ def mikkola_coe_vf(k, p, ecc, inc, raan, argp, nu, tof): @hjit("Tuple([V,V])(f,V,V,f)") -def mikkola_rv_hf(k, r0, v0, tof): # rtol=None +def mikkola_rv_hf(k, r0, v0, tof): """Raw algorithm for Mikkola's Kepler solver. Parameters @@ -129,8 +129,6 @@ def mikkola_rv_hf(k, r0, v0, tof): # rtol=None Velocity vector. tof : float Time of flight. - rtol : float - This method does not require tolerance since it is non-iterative. Returns ------- @@ -150,7 +148,7 @@ def mikkola_rv_hf(k, r0, v0, tof): # rtol=None @gjit("void(f,f[:],f[:],f,f[:],f[:])", "(),(n),(n),()->(n),(n)") -def mikkola_rv_gf(k, r0, v0, tof, rr, vv): # rtol=None +def mikkola_rv_gf(k, r0, v0, tof, rr, vv): """ Vectorized mikkola_rv """ From c6d2bbbb80aa2ba9e83fe0bb35babc6bb911d506 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 8 Jan 2024 16:24:32 +0100 Subject: [PATCH 073/346] jit pimienta --- src/hapsira/core/propagation/__init__.py | 6 +- src/hapsira/core/propagation/pimienta.py | 61 +++++++++++++++------ src/hapsira/twobody/propagation/pimienta.py | 4 +- tests/tests_core/test_core_propagation.py | 6 +- 4 files changed, 52 insertions(+), 25 deletions(-) diff --git a/src/hapsira/core/propagation/__init__.py b/src/hapsira/core/propagation/__init__.py index 31f016a39..4e1e7f88e 100644 --- a/src/hapsira/core/propagation/__init__.py +++ b/src/hapsira/core/propagation/__init__.py @@ -11,7 +11,7 @@ # from hapsira.core.propagation.gooding import gooding, gooding_coe # from hapsira.core.propagation.markley import markley, markley_coe # from hapsira.core.propagation.mikkola import mikkola, mikkola_coe -from hapsira.core.propagation.pimienta import pimienta, pimienta_coe +# from hapsira.core.propagation.pimienta import pimienta, pimienta_coe from hapsira.core.propagation.recseries import recseries, recseries_coe from hapsira.core.propagation.vallado import vallado @@ -25,8 +25,8 @@ # "mikkola", # "markley_coe", # "markley", - "pimienta_coe", - "pimienta", + # "pimienta_coe", + # "pimienta", # "gooding_coe", # "gooding", # "danby_coe", diff --git a/src/hapsira/core/propagation/pimienta.py b/src/hapsira/core/propagation/pimienta.py index f468fc07b..dfc471aa6 100644 --- a/src/hapsira/core/propagation/pimienta.py +++ b/src/hapsira/core/propagation/pimienta.py @@ -1,17 +1,28 @@ -from numba import njit as jit -import numpy as np +from math import sqrt -from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf -from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL -from ..jit import array_to_V_hf +from ..angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf +from ..elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf, hjit, vjit, gjit -@jit -def pimienta_coe(k, p, ecc, inc, raan, argp, nu, tof): +__all__ = [ + "pimienta_coe_hf", + "pimienta_coe_vf", + "pimienta_rv_hf", + "pimienta_rv_gf", +] + + +@hjit("f(f,f,f,f,f,f,f,f)") +def pimienta_coe_hf(k, p, ecc, inc, raan, argp, nu, tof): + """ + Scalar pimienta_coe + """ + q = p / (1 + ecc) # TODO: Do something to allow parabolic and hyperbolic orbits? - n = np.sqrt(k * (1 - ecc) ** 3 / q**3) + n = sqrt(k * (1 - ecc) ** 3 / q**3) M0 = E_to_M_hf(nu_to_E_hf(nu, ecc), ecc) M = M0 + n * tof @@ -20,7 +31,7 @@ def pimienta_coe(k, p, ecc, inc, raan, argp, nu, tof): c3 = 5 / 2 + 560 * ecc a = 15 * (1 - ecc) / c3 b = -M / c3 - y = np.sqrt(b**2 / 4 + a**3 / 27) + y = sqrt(b**2 / 4 + a**3 / 27) # Equation (33) x_bar = (-b / 2 + y) ** (1 / 3) - (b / 2 + y) ** (1 / 3) @@ -337,8 +348,17 @@ def pimienta_coe(k, p, ecc, inc, raan, argp, nu, tof): return E_to_nu_hf(E, ecc) -@jit -def pimienta(k, r0, v0, tof): +@vjit("f(f,f,f,f,f,f,f,f)") +def pimienta_coe_vf(k, p, ecc, inc, raan, argp, nu, tof): + """ + Vectorized pimienta_coe + """ + + return pimienta_coe_hf(k, p, ecc, inc, raan, argp, nu, tof) + + +@hjit("Tuple([V,V])(f,V,V,f)") +def pimienta_rv_hf(k, r0, v0, tof): """Raw algorithm for Adonis' Pimienta and John L. Crassidis 15th order polynomial Kepler solver. @@ -369,9 +389,18 @@ def pimienta(k, r0, v0, tof): # TODO: implement hyperbolic case # Solve first for eccentricity and mean anomaly - p, ecc, inc, raan, argp, nu = rv2coe_hf( - k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL - ) - nu = pimienta_coe(k, p, ecc, inc, raan, argp, nu, tof) + p, ecc, inc, raan, argp, nu = rv2coe_hf(k, r0, v0, RV2COE_TOL) + nu = pimienta_coe_hf(k, p, ecc, inc, raan, argp, nu, tof) + + return coe2rv_hf(k, p, ecc, inc, raan, argp, nu) + - return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) +@gjit("void(f,f[:],f[:],f,f[:],f[:])", "(),(n),(n),()->(n),(n)") +def pimienta_rv_gf(k, r0, v0, tof, rr, vv): + """ + Vectorized pimienta_rv + """ + + (rr[0], rr[1], rr[2]), (vv[0], vv[1], vv[2]) = pimienta_rv_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), tof + ) diff --git a/src/hapsira/twobody/propagation/pimienta.py b/src/hapsira/twobody/propagation/pimienta.py index 0e21908db..0a61b1735 100644 --- a/src/hapsira/twobody/propagation/pimienta.py +++ b/src/hapsira/twobody/propagation/pimienta.py @@ -2,7 +2,7 @@ from astropy import units as u -from hapsira.core.propagation import pimienta_coe as pimienta_fast +from hapsira.core.propagation.pimienta import pimienta_coe_vf from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import ClassicalState @@ -30,7 +30,7 @@ def propagate(self, state, tof): state = state.to_classical() nu = ( - pimienta_fast( + pimienta_coe_vf( state.attractor.k.to_value(u.km**3 / u.s**2), *state.to_value(), tof.to_value(u.s), diff --git a/tests/tests_core/test_core_propagation.py b/tests/tests_core/test_core_propagation.py index 114c6c27d..5be66e4bf 100644 --- a/tests/tests_core/test_core_propagation.py +++ b/tests/tests_core/test_core_propagation.py @@ -2,9 +2,6 @@ from astropy.tests.helper import assert_quantity_allclose import pytest -from hapsira.core.propagation import ( - pimienta_coe, -) from hapsira.core.propagation.danby import danby_coe_vf, DANBY_NUMITER, DANBY_RTOL from hapsira.core.propagation.farnocchia import farnocchia_coe_vf from hapsira.core.propagation.gooding import ( @@ -14,6 +11,7 @@ ) from hapsira.core.propagation.markley import markley_coe_vf from hapsira.core.propagation.mikkola import mikkola_coe_vf +from hapsira.core.propagation.pimienta import pimienta_coe_vf from hapsira.examples import iss @@ -22,7 +20,7 @@ [ lambda *args: danby_coe_vf(*args, DANBY_NUMITER, DANBY_RTOL), markley_coe_vf, - pimienta_coe, + pimienta_coe_vf, mikkola_coe_vf, farnocchia_coe_vf, lambda *args: gooding_coe_vf(*args, GOODING_NUMITER, GOODING_RTOL), From 6bc04da415ed46ce7df306100fdc6479a73340ff Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 8 Jan 2024 16:58:35 +0100 Subject: [PATCH 074/346] jit recseries --- src/hapsira/core/propagation/__init__.py | 6 +- src/hapsira/core/propagation/recseries.py | 115 +++++++++++++++---- src/hapsira/twobody/propagation/recseries.py | 26 +++-- tests/tests_core/test_core_propagation.py | 14 +++ 4 files changed, 124 insertions(+), 37 deletions(-) diff --git a/src/hapsira/core/propagation/__init__.py b/src/hapsira/core/propagation/__init__.py index 4e1e7f88e..40d1bae13 100644 --- a/src/hapsira/core/propagation/__init__.py +++ b/src/hapsira/core/propagation/__init__.py @@ -12,7 +12,7 @@ # from hapsira.core.propagation.markley import markley, markley_coe # from hapsira.core.propagation.mikkola import mikkola, mikkola_coe # from hapsira.core.propagation.pimienta import pimienta, pimienta_coe -from hapsira.core.propagation.recseries import recseries, recseries_coe +# from hapsira.core.propagation.recseries import recseries, recseries_coe from hapsira.core.propagation.vallado import vallado __all__ = [ @@ -31,6 +31,6 @@ # "gooding", # "danby_coe", # "danby", - "recseries_coe", - "recseries", + # "recseries_coe", + # "recseries", ] diff --git a/src/hapsira/core/propagation/recseries.py b/src/hapsira/core/propagation/recseries.py index 99ff501d3..44ef6a92e 100644 --- a/src/hapsira/core/propagation/recseries.py +++ b/src/hapsira/core/propagation/recseries.py @@ -1,13 +1,32 @@ -from numba import njit as jit -import numpy as np +from math import floor, pi, sin, sqrt -from hapsira.core.angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf -from hapsira.core.elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL -from ..jit import array_to_V_hf +from ..angles import E_to_M_hf, E_to_nu_hf, nu_to_E_hf +from ..elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..jit import array_to_V_hf, hjit, vjit, gjit -@jit -def recseries_coe( +__all__ = [ + "recseries_coe_hf", + "recseries_coe_vf", + "recseries_rv_hf", + "recseries_rv_gf", + "RECSERIES_METHOD_RTOL", + "RECSERIES_METHOD_ORDER", + "RECSERIES_ORDER", + "RECSERIES_NUMITER", + "RECSERIES_RTOL", +] + + +RECSERIES_METHOD_RTOL = 0 +RECSERIES_METHOD_ORDER = 1 +RECSERIES_ORDER = 8 +RECSERIES_NUMITER = 100 +RECSERIES_RTOL = 1e-8 + + +@hjit("f(f,f,f,f,f,f,f,f,i8,i8,i8,f)") +def recseries_coe_hf( k, p, ecc, @@ -16,15 +35,19 @@ def recseries_coe( argp, nu, tof, - method="rtol", - order=8, - numiter=100, - rtol=1e-8, + method, + order, + numiter, + rtol, ): + """ + Scalar recseries_coe + """ + # semi-major axis semi_axis_a = p / (1 - ecc**2) # mean angular motion - n = np.sqrt(k / np.abs(semi_axis_a) ** 3) + n = sqrt(k / abs(semi_axis_a) ** 3) if ecc == 0: # Solving for circular orbit @@ -34,7 +57,7 @@ def recseries_coe( # final mean anaomaly M = M0 + n * tof # snapping anomaly to [0,pi] range - nu = M - 2 * np.pi * np.floor(M / 2 / np.pi) + nu = M - 2 * pi * floor(M / 2 / pi) return nu @@ -46,12 +69,12 @@ def recseries_coe( # final mean anaomaly M = M0 + n * tof # snapping anomaly to [0,pi] range - M = M - 2 * np.pi * np.floor(M / 2 / np.pi) + M = M - 2 * pi * floor(M / 2 / pi) # set recursion iteration - if method == "rtol": + if method == RECSERIES_METHOD_RTOL: Niter = numiter - elif method == "order": + elif method == RECSERIES_METHOD_ORDER: Niter = order else: raise ValueError("Unknown recursion termination method ('rtol','order').") @@ -59,7 +82,7 @@ def recseries_coe( # compute eccentric anomaly through recursive series E = M + ecc # Using initial guess from vallado to improve convergence for i in range(0, Niter): - En = M + ecc * np.sin(E) + En = M + ecc * sin(E) # check for break condition if method == "rtol" and (abs(En - E) / abs(E)) < rtol: break @@ -74,8 +97,43 @@ def recseries_coe( return nu -@jit -def recseries(k, r0, v0, tof, method="rtol", order=8, numiter=100, rtol=1e-8): +@vjit("f(f,f,f,f,f,f,f,f,i8,i8,i8,f)") +def recseries_coe_vf( + k, + p, + ecc, + inc, + raan, + argp, + nu, + tof, + method, + order, + numiter, + rtol, +): + """ + Vectorized recseries_coe + """ + + return recseries_coe_hf( + k, + p, + ecc, + inc, + raan, + argp, + nu, + tof, + method, + order, + numiter, + rtol, + ) + + +@hjit("Tuple([V,V])(f,V,V,f,i8,i8,i8,f)") +def recseries_rv_hf(k, r0, v0, tof, method, order, numiter, rtol): """Kepler solver for elliptical orbits with recursive series approximation method. The order of the series is a user defined parameter. @@ -113,11 +171,20 @@ def recseries(k, r0, v0, tof, method="rtol", order=8, numiter=100, rtol=1e-8): with DOI: http://dx.doi.org/10.13140/RG.2.2.18578.58563/1 """ # Solve first for eccentricity and mean anomaly - p, ecc, inc, raan, argp, nu = rv2coe_hf( - k, array_to_V_hf(r0), array_to_V_hf(v0), RV2COE_TOL - ) - nu = recseries_coe( + p, ecc, inc, raan, argp, nu = rv2coe_hf(k, r0, v0, RV2COE_TOL) + nu = recseries_coe_hf( k, p, ecc, inc, raan, argp, nu, tof, method, order, numiter, rtol ) - return np.array(coe2rv_hf(k, p, ecc, inc, raan, argp, nu)) + return coe2rv_hf(k, p, ecc, inc, raan, argp, nu) + + +@gjit("void(f,f[:],f[:],f,i8,i8,i8,f,f[:],f[:])", "(),(n),(n),(),(),(),(),()->(n),(n)") +def recseries_rv_gf(k, r0, v0, tof, method, order, numiter, rtol, rr, vv): + """ + Vectorized recseries_rv + """ + + (rr[0], rr[1], rr[2]), (vv[0], vv[1], vv[2]) = recseries_rv_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), tof, method, order, numiter, rtol + ) diff --git a/src/hapsira/twobody/propagation/recseries.py b/src/hapsira/twobody/propagation/recseries.py index cff5a0c23..5536152b8 100644 --- a/src/hapsira/twobody/propagation/recseries.py +++ b/src/hapsira/twobody/propagation/recseries.py @@ -2,7 +2,13 @@ from astropy import units as u -from hapsira.core.propagation import recseries_coe as recseries_fast +from hapsira.core.propagation.recseries import ( + recseries_coe_vf, + RECSERIES_METHOD_RTOL, + RECSERIES_ORDER, + RECSERIES_NUMITER, + RECSERIES_RTOL, +) from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import ClassicalState @@ -29,10 +35,10 @@ class RecseriesPropagator: def __init__( self, - method="rtol", - order=8, - numiter=100, - rtol=1e-8, + method=RECSERIES_METHOD_RTOL, + order=RECSERIES_ORDER, + numiter=RECSERIES_NUMITER, + rtol=RECSERIES_RTOL, ): self._method = method self._order = order @@ -43,14 +49,14 @@ def propagate(self, state, tof): state = state.to_classical() nu = ( - recseries_fast( + recseries_coe_vf( state.attractor.k.to_value(u.km**3 / u.s**2), *state.to_value(), tof.to_value(u.s), - method=self._method, - order=self._order, - numiter=self._numiter, - rtol=self._rtol, + self._method, + self._order, + self._numiter, + self._rtol, ) << u.rad ) diff --git a/tests/tests_core/test_core_propagation.py b/tests/tests_core/test_core_propagation.py index 5be66e4bf..1f4fa160a 100644 --- a/tests/tests_core/test_core_propagation.py +++ b/tests/tests_core/test_core_propagation.py @@ -12,6 +12,13 @@ from hapsira.core.propagation.markley import markley_coe_vf from hapsira.core.propagation.mikkola import mikkola_coe_vf from hapsira.core.propagation.pimienta import pimienta_coe_vf +from hapsira.core.propagation.recseries import ( + recseries_coe_vf, + RECSERIES_METHOD_RTOL, + RECSERIES_ORDER, + RECSERIES_NUMITER, + RECSERIES_RTOL, +) from hapsira.examples import iss @@ -24,6 +31,13 @@ mikkola_coe_vf, farnocchia_coe_vf, lambda *args: gooding_coe_vf(*args, GOODING_NUMITER, GOODING_RTOL), + lambda *args: recseries_coe_vf( + *args, + RECSERIES_METHOD_RTOL, + RECSERIES_ORDER, + RECSERIES_NUMITER, + RECSERIES_RTOL, + ), ], ) def test_propagate_with_coe(propagator_coe): From c99f7847bd2913fde59899de674e3222ab936821 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 8 Jan 2024 17:33:22 +0100 Subject: [PATCH 075/346] jit vallado --- src/hapsira/core/propagation/__init__.py | 4 +- src/hapsira/core/propagation/vallado.py | 96 ++++++++++++++++++---- src/hapsira/twobody/propagation/vallado.py | 24 +----- tests/tests_core/test_core_propagation.py | 2 + 4 files changed, 87 insertions(+), 39 deletions(-) diff --git a/src/hapsira/core/propagation/__init__.py b/src/hapsira/core/propagation/__init__.py index 40d1bae13..fb4bc643d 100644 --- a/src/hapsira/core/propagation/__init__.py +++ b/src/hapsira/core/propagation/__init__.py @@ -13,14 +13,14 @@ # from hapsira.core.propagation.mikkola import mikkola, mikkola_coe # from hapsira.core.propagation.pimienta import pimienta, pimienta_coe # from hapsira.core.propagation.recseries import recseries, recseries_coe -from hapsira.core.propagation.vallado import vallado +# from hapsira.core.propagation.vallado import vallado __all__ = [ "cowell", "func_twobody", # "farnocchia_coe", # "farnocchia", - "vallado", + # "vallado", # "mikkola_coe", # "mikkola", # "markley_coe", diff --git a/src/hapsira/core/propagation/vallado.py b/src/hapsira/core/propagation/vallado.py index ce94dcf77..076892e31 100644 --- a/src/hapsira/core/propagation/vallado.py +++ b/src/hapsira/core/propagation/vallado.py @@ -1,13 +1,25 @@ -from numba import njit as jit -import numpy as np +from math import log, sqrt -from ..jit import array_to_V_hf -from ..math.linalg import norm_hf +from ..elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL +from ..math.linalg import add_VV_hf, matmul_VV_hf, mul_Vs_hf, norm_hf, sign_hf from ..math.special import stumpff_c2_hf, stumpff_c3_hf +from ..jit import array_to_V_hf, hjit, vjit, gjit -@jit -def vallado(k, r0, v0, tof, numiter): +__all__ = [ + "vallado_coe_hf", + "vallado_coe_vf", + "vallado_rv_hf", + "vallado_rv_gf", + "VALLADO_NUMITER", +] + + +VALLADO_NUMITER = 350 + + +@hjit("Tuple([f,f,f,f])(f,V,V,f,i8)") +def _vallado_hf(k, r0, v0, tof, numiter): r"""Solves Kepler's Equation by applying a Newton-Raphson method. If the position of a body along its orbit wants to be computed @@ -46,9 +58,9 @@ def vallado(k, r0, v0, tof, numiter): ---------- k : float Standard gravitational parameter. - r0 : numpy.ndarray + r0 : tuple[float,float,float] Initial position vector. - v0 : numpy.ndarray + v0 : tuple[float,float,float] Initial velocity vector. tof : float Time of flight. @@ -73,10 +85,10 @@ def vallado(k, r0, v0, tof, numiter): """ # Cache some results - dot_r0v0 = r0 @ v0 - norm_r0 = norm_hf(array_to_V_hf(r0)) + dot_r0v0 = matmul_VV_hf(r0, v0) + norm_r0 = norm_hf(r0) sqrt_mu = k**0.5 - alpha = -(v0 @ v0) / k + 2 / norm_r0 + alpha = -matmul_VV_hf(v0, v0) / k + 2 / norm_r0 # First guess if alpha > 0: @@ -85,14 +97,11 @@ def vallado(k, r0, v0, tof, numiter): elif alpha < 0: # Hyperbolic orbit xi_new = ( - np.sign(tof) + sign_hf(tof) * (-1 / alpha) ** 0.5 - * np.log( + * log( (-2 * k * alpha * tof) - / ( - dot_r0v0 - + np.sign(tof) * np.sqrt(-k / alpha) * (1 - norm_r0 * alpha) - ) + / (dot_r0v0 + sign_hf(tof) * sqrt(-k / alpha) * (1 - norm_r0 * alpha)) ) ) else: @@ -137,3 +146,56 @@ def vallado(k, r0, v0, tof, numiter): fdot = sqrt_mu / (norm_r * norm_r0) * xi * (psi * c3_psi - 1) return f, g, fdot, gdot + + +@hjit("Tuple([V,V])(f,V,V,f,i8)") +def vallado_rv_hf(k, r0, v0, tof, numiter): + """ + Scalar vallado_rv + """ + + # Compute Lagrange coefficients + f, g, fdot, gdot = _vallado_hf(k, r0, v0, tof, numiter) + + assert ( + abs(f * gdot - fdot * g - 1) < 1e-5 + ), "Internal error, solution is not consistent" # Fixed tolerance + + # Return position and velocity vectors + r = add_VV_hf(mul_Vs_hf(r0, f), mul_Vs_hf(v0, g)) + v = add_VV_hf(mul_Vs_hf(r0, fdot), mul_Vs_hf(v0, gdot)) + + return r, v + + +@gjit("void(f,f[:],f[:],f,i8,f[:],f[:])", "(),(n),(n),(),()->(n),(n)") +def vallado_rv_gf(k, r0, v0, tof, numiter, rr, vv): + """ + Vectorized vallado_rv + """ + + (rr[0], rr[1], rr[2]), (vv[0], vv[1], vv[2]) = vallado_rv_hf( + k, array_to_V_hf(r0), array_to_V_hf(v0), tof, numiter + ) + + +@hjit("f(f,f,f,f,f,f,f,f,i8)") +def vallado_coe_hf(k, p, ecc, inc, raan, argp, nu, tof, numiter): + """ + Scalar vallado_coe + """ + + r0, v0 = coe2rv_hf(k, p, ecc, inc, raan, argp, nu) + rr, vv = vallado_rv_hf(k, r0, v0, tof, numiter) + _, _, _, _, _, nu_ = rv2coe_hf(k, rr, vv, RV2COE_TOL) + + return nu_ + + +@vjit("f(f,f,f,f,f,f,f,f,i8)") +def vallado_coe_vf(k, p, ecc, inc, raan, argp, nu, tof, numiter): + """ + Vectorized vallado_coe + """ + + return vallado_coe_hf(k, p, ecc, inc, raan, argp, nu, tof, numiter) diff --git a/src/hapsira/twobody/propagation/vallado.py b/src/hapsira/twobody/propagation/vallado.py index 06d203929..2fb197117 100644 --- a/src/hapsira/twobody/propagation/vallado.py +++ b/src/hapsira/twobody/propagation/vallado.py @@ -1,9 +1,8 @@ import sys from astropy import units as u -import numpy as np -from hapsira.core.propagation import vallado as vallado_fast +from hapsira.core.propagation.vallado import vallado_rv_gf, VALLADO_NUMITER from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import RVState @@ -12,21 +11,6 @@ sys.modules[__name__].__class__ = OldPropagatorModule -def vallado(k, r0, v0, tof, *, numiter): - # Compute Lagrange coefficients - f, g, fdot, gdot = vallado_fast(k, r0, v0, tof, numiter) - - assert ( - np.abs(f * gdot - fdot * g - 1) < 1e-5 - ), "Internal error, solution is not consistent" # Fixed tolerance - - # Return position and velocity vectors - r = f * r0 + g * v0 - v = fdot * r0 + gdot * v0 - - return r, v - - class ValladoPropagator: """Propagates Keplerian orbit using Vallado's method. @@ -43,17 +27,17 @@ class ValladoPropagator: PropagatorKind.ELLIPTIC | PropagatorKind.PARABOLIC | PropagatorKind.HYPERBOLIC ) - def __init__(self, numiter=350): + def __init__(self, numiter=VALLADO_NUMITER): self._numiter = numiter def propagate(self, state, tof): state = state.to_vectors() - r_raw, v_raw = vallado( + r_raw, v_raw = vallado_rv_gf( state.attractor.k.to_value(u.km**3 / u.s**2), *state.to_value(), tof.to_value(u.s), - numiter=self._numiter, + self._numiter, ) r = r_raw << u.km v = v_raw << (u.km / u.s) diff --git a/tests/tests_core/test_core_propagation.py b/tests/tests_core/test_core_propagation.py index 1f4fa160a..8dc4a3711 100644 --- a/tests/tests_core/test_core_propagation.py +++ b/tests/tests_core/test_core_propagation.py @@ -19,6 +19,7 @@ RECSERIES_NUMITER, RECSERIES_RTOL, ) +from hapsira.core.propagation.vallado import vallado_coe_vf, VALLADO_NUMITER from hapsira.examples import iss @@ -38,6 +39,7 @@ RECSERIES_NUMITER, RECSERIES_RTOL, ), + lambda *args: vallado_coe_vf(*args, VALLADO_NUMITER), ], ) def test_propagate_with_coe(propagator_coe): From be0ae89dea496ce30e78783502d8a2880cc2ffda Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 3 Jan 2024 14:32:15 +0100 Subject: [PATCH 076/346] ivp into fld --- src/hapsira/core/math/ivp/__init__.py | 0 src/hapsira/core/math/{ivp.py => ivp/_api.py} | 0 2 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 src/hapsira/core/math/ivp/__init__.py rename src/hapsira/core/math/{ivp.py => ivp/_api.py} (100%) diff --git a/src/hapsira/core/math/ivp/__init__.py b/src/hapsira/core/math/ivp/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/hapsira/core/math/ivp.py b/src/hapsira/core/math/ivp/_api.py similarity index 100% rename from src/hapsira/core/math/ivp.py rename to src/hapsira/core/math/ivp/_api.py From 7fbcb06273ad226232577b684154986f5a653d24 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 3 Jan 2024 15:05:20 +0100 Subject: [PATCH 077/346] isolate dop8853 --- contrib/CR3BP/CR3BP.py | 4 +--- src/hapsira/core/math/ivp/__init__.py | 5 +++++ src/hapsira/core/math/ivp/_api.py | 17 ++++++++++++----- src/hapsira/core/propagation/cowell.py | 3 +-- 4 files changed, 19 insertions(+), 10 deletions(-) diff --git a/contrib/CR3BP/CR3BP.py b/contrib/CR3BP/CR3BP.py index 237f8e62f..d581e4683 100644 --- a/contrib/CR3BP/CR3BP.py +++ b/contrib/CR3BP/CR3BP.py @@ -32,7 +32,7 @@ from numba import njit as jit import numpy as np -from hapsira.core.math.ivp import DOP853, solve_ivp +from hapsira.core.math.ivp import solve_ivp @jit @@ -314,7 +314,6 @@ def propagate(mu, r0, v0, tofs, rtol=1e-11, f=func_CR3BP): args=(mu,), rtol=rtol, atol=1e-12, - method=DOP853, dense_output=True, ) @@ -371,7 +370,6 @@ def propagateSTM(mu, r0, v0, STM0, tofs, rtol=1e-11, f=func_CR3BP_STM): args=(mu,), rtol=rtol, atol=1e-12, - method=DOP853, dense_output=True, ) diff --git a/src/hapsira/core/math/ivp/__init__.py b/src/hapsira/core/math/ivp/__init__.py index e69de29bb..da2192fdc 100644 --- a/src/hapsira/core/math/ivp/__init__.py +++ b/src/hapsira/core/math/ivp/__init__.py @@ -0,0 +1,5 @@ +from ._api import solve_ivp + +__all__ = [ + "solve_ivp", +] diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 29e4f7674..f92af23dd 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -1,6 +1,13 @@ -from scipy.integrate import DOP853, solve_ivp +from scipy.integrate import DOP853, solve_ivp as solve_ivp_ -__all__ = [ - "DOP853", - "solve_ivp", -] + +def solve_ivp(*args, **kwargs): + """ + Wrapper for `scipy.integrate.solve_ivp` + """ + + return solve_ivp_( + *args, + method=DOP853, + **kwargs, + ) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 8c8bae666..c0959cc15 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -2,7 +2,7 @@ from hapsira.core.propagation.base import func_twobody -from ..math.ivp import DOP853, solve_ivp +from ..math.ivp import solve_ivp def cowell(k, r, v, tofs, rtol=1e-11, *, events=None, f=func_twobody): @@ -18,7 +18,6 @@ def cowell(k, r, v, tofs, rtol=1e-11, *, events=None, f=func_twobody): args=(k,), rtol=rtol, atol=1e-12, - method=DOP853, dense_output=True, events=events, ) From 72e817eb79399d0ac0aa37eaffaa6c275720649e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 8 Jan 2024 13:37:13 +0100 Subject: [PATCH 078/346] draft --- src/hapsira/core/math/ivp/_api.py | 773 +++++++++++++++++++++++++++++- src/hapsira/core/math/ivp/_rk.py | 0 2 files changed, 766 insertions(+), 7 deletions(-) create mode 100644 src/hapsira/core/math/ivp/_rk.py diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index f92af23dd..5dc35a123 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -1,13 +1,772 @@ -from scipy.integrate import DOP853, solve_ivp as solve_ivp_ +from scipy.integrate import DOP853 +import numpy as np -def solve_ivp(*args, **kwargs): + +class OdeResult(dict): + """Represents the optimization result. + + Attributes + ---------- + x : ndarray + The solution of the optimization. + success : bool + Whether or not the optimizer exited successfully. + status : int + Termination status of the optimizer. Its value depends on the + underlying solver. Refer to `message` for details. + message : str + Description of the cause of the termination. + fun, jac, hess: ndarray + Values of objective function, its Jacobian and its Hessian (if + available). The Hessians may be approximations, see the documentation + of the function in question. + hess_inv : object + Inverse of the objective function's Hessian; may be an approximation. + Not available for all solvers. The type of this attribute may be + either np.ndarray or scipy.sparse.linalg.LinearOperator. + nfev, njev, nhev : int + Number of evaluations of the objective functions and of its + Jacobian and Hessian. + nit : int + Number of iterations performed by the optimizer. + maxcv : float + The maximum constraint violation. + + Notes + ----- + Depending on the specific solver being used, `OptimizeResult` may + not have all attributes listed here, and they may have additional + attributes not listed here. Since this class is essentially a + subclass of dict with attribute accessors, one can see which + attributes are available using the `OptimizeResult.keys` method. + """ + + def __getattr__(self, name): + try: + return self[name] + except KeyError as e: + raise AttributeError(name) from e + + __setattr__ = dict.__setitem__ + __delattr__ = dict.__delitem__ + + # def __repr__(self): + # order_keys = ['message', 'success', 'status', 'fun', 'funl', 'x', 'xl', + # 'col_ind', 'nit', 'lower', 'upper', 'eqlin', 'ineqlin', + # 'converged', 'flag', 'function_calls', 'iterations', + # 'root'] + # order_keys = getattr(self, '_order_keys', order_keys) + # # 'slack', 'con' are redundant with residuals + # # 'crossover_nit' is probably not interesting to most users + # omit_keys = {'slack', 'con', 'crossover_nit', '_order_keys'} + + # def key(item): + # try: + # return order_keys.index(item[0].lower()) + # except ValueError: # item not in list + # return np.inf + + # def omit_redundant(items): + # for item in items: + # if item[0] in omit_keys: + # continue + # yield item + + # def item_sorter(d): + # return sorted(omit_redundant(d.items()), key=key) + + # if self.keys(): + # return _dict_formatter(self, sorter=item_sorter) + # else: + # return self.__class__.__name__ + "()" + + def __dir__(self): + return list(self.keys()) + + +def solve_event_equation(event, sol, t_old, t): + """Solve an equation corresponding to an ODE event. + + The equation is ``event(t, y(t)) = 0``, here ``y(t)`` is known from an + ODE solver using some sort of interpolation. It is solved by + `scipy.optimize.brentq` with xtol=atol=4*EPS. + + Parameters + ---------- + event : callable + Function ``event(t, y)``. + sol : callable + Function ``sol(t)`` which evaluates an ODE solution between `t_old` + and `t`. + t_old, t : float + Previous and new values of time. They will be used as a bracketing + interval. + + Returns + ------- + root : float + Found solution. + """ + from scipy.optimize import brentq + + return brentq(lambda t: event(t, sol(t)), t_old, t, xtol=4 * EPS, rtol=4 * EPS) + # https://github.com/scipy/scipy/blob/ea4d1f1330950bee74396e427fe6330424907621/scipy/optimize/_zeros_py.py#L682 + # https://github.com/scipy/scipy/blob/ea4d1f1330950bee74396e427fe6330424907621/scipy/optimize/Zeros/brentq.c#L37 + + +def handle_events(sol, events, active_events, is_terminal, t_old, t): + """Helper function to handle events. + + Parameters + ---------- + sol : DenseOutput + Function ``sol(t)`` which evaluates an ODE solution between `t_old` + and `t`. + events : list of callables, length n_events + Event functions with signatures ``event(t, y)``. + active_events : ndarray + Indices of events which occurred. + is_terminal : ndarray, shape (n_events,) + Which events are terminal. + t_old, t : float + Previous and new values of time. + + Returns + ------- + root_indices : ndarray + Indices of events which take zero between `t_old` and `t` and before + a possible termination. + roots : ndarray + Values of t at which events occurred. + terminate : bool + Whether a terminal event occurred. + """ + roots = [ + solve_event_equation(events[event_index], sol, t_old, t) + for event_index in active_events + ] + + roots = np.asarray(roots) + + if np.any(is_terminal[active_events]): + if t > t_old: + order = np.argsort(roots) + else: + order = np.argsort(-roots) + active_events = active_events[order] + roots = roots[order] + t = np.nonzero(is_terminal[active_events])[0][0] + active_events = active_events[: t + 1] + roots = roots[: t + 1] + terminate = True + else: + terminate = False + + return active_events, roots, terminate + + +def prepare_events(events): + """Standardize event functions and extract is_terminal and direction.""" + if callable(events): + events = (events,) + + if events is not None: + is_terminal = np.empty(len(events), dtype=bool) + direction = np.empty(len(events)) + for i, event in enumerate(events): + try: + is_terminal[i] = event.terminal + except AttributeError: + is_terminal[i] = False + + try: + direction[i] = event.direction + except AttributeError: + direction[i] = 0 + else: + is_terminal = None + direction = None + + return events, is_terminal, direction + + +def find_active_events(g, g_new, direction): + """Find which event occurred during an integration step. + + Parameters + ---------- + g, g_new : array_like, shape (n_events,) + Values of event functions at a current and next points. + direction : ndarray, shape (n_events,) + Event "direction" according to the definition in `solve_ivp`. + + Returns + ------- + active_events : ndarray + Indices of events which occurred during the step. """ - Wrapper for `scipy.integrate.solve_ivp` + g, g_new = np.asarray(g), np.asarray(g_new) + up = (g <= 0) & (g_new >= 0) + down = (g >= 0) & (g_new <= 0) + either = up | down + mask = up & (direction > 0) | down & (direction < 0) | either & (direction == 0) + + return np.nonzero(mask)[0] + + +def solve_ivp( + fun, + t_span, + y0, + method=DOP853, + t_eval=None, + dense_output=False, + events=None, + vectorized=False, + args=None, + **options, +): + """Solve an initial value problem for a system of ODEs. + + This function numerically integrates a system of ordinary differential + equations given an initial value:: + + dy / dt = f(t, y) + y(t0) = y0 + + Here t is a 1-D independent variable (time), y(t) is an + N-D vector-valued function (state), and an N-D + vector-valued function f(t, y) determines the differential equations. + The goal is to find y(t) approximately satisfying the differential + equations, given an initial value y(t0)=y0. + + Some of the solvers support integration in the complex domain, but note + that for stiff ODE solvers, the right-hand side must be + complex-differentiable (satisfy Cauchy-Riemann equations [11]_). + To solve a problem in the complex domain, pass y0 with a complex data type. + Another option always available is to rewrite your problem for real and + imaginary parts separately. + + Parameters + ---------- + fun : callable + Right-hand side of the system: the time derivative of the state ``y`` + at time ``t``. The calling signature is ``fun(t, y)``, where ``t`` is a + scalar and ``y`` is an ndarray with ``len(y) = len(y0)``. ``fun`` must + return an array of the same shape as ``y``. See `vectorized` for more + information. + t_span : 2-member sequence + Interval of integration (t0, tf). The solver starts with t=t0 and + integrates until it reaches t=tf. Both t0 and tf must be floats + or values interpretable by the float conversion function. + y0 : array_like, shape (n,) + Initial state. For problems in the complex domain, pass `y0` with a + complex data type (even if the initial value is purely real). + method : string or `OdeSolver`, optional + Integration method to use: + + * 'RK45' (default): Explicit Runge-Kutta method of order 5(4) [1]_. + The error is controlled assuming accuracy of the fourth-order + method, but steps are taken using the fifth-order accurate + formula (local extrapolation is done). A quartic interpolation + polynomial is used for the dense output [2]_. Can be applied in + the complex domain. + * 'RK23': Explicit Runge-Kutta method of order 3(2) [3]_. The error + is controlled assuming accuracy of the second-order method, but + steps are taken using the third-order accurate formula (local + extrapolation is done). A cubic Hermite polynomial is used for the + dense output. Can be applied in the complex domain. + * 'DOP853': Explicit Runge-Kutta method of order 8 [13]_. + Python implementation of the "DOP853" algorithm originally + written in Fortran [14]_. A 7-th order interpolation polynomial + accurate to 7-th order is used for the dense output. + Can be applied in the complex domain. + * 'Radau': Implicit Runge-Kutta method of the Radau IIA family of + order 5 [4]_. The error is controlled with a third-order accurate + embedded formula. A cubic polynomial which satisfies the + collocation conditions is used for the dense output. + * 'BDF': Implicit multi-step variable-order (1 to 5) method based + on a backward differentiation formula for the derivative + approximation [5]_. The implementation follows the one described + in [6]_. A quasi-constant step scheme is used and accuracy is + enhanced using the NDF modification. Can be applied in the + complex domain. + * 'LSODA': Adams/BDF method with automatic stiffness detection and + switching [7]_, [8]_. This is a wrapper of the Fortran solver + from ODEPACK. + + Explicit Runge-Kutta methods ('RK23', 'RK45', 'DOP853') should be used + for non-stiff problems and implicit methods ('Radau', 'BDF') for + stiff problems [9]_. Among Runge-Kutta methods, 'DOP853' is recommended + for solving with high precision (low values of `rtol` and `atol`). + + If not sure, first try to run 'RK45'. If it makes unusually many + iterations, diverges, or fails, your problem is likely to be stiff and + you should use 'Radau' or 'BDF'. 'LSODA' can also be a good universal + choice, but it might be somewhat less convenient to work with as it + wraps old Fortran code. + + You can also pass an arbitrary class derived from `OdeSolver` which + implements the solver. + t_eval : array_like or None, optional + Times at which to store the computed solution, must be sorted and lie + within `t_span`. If None (default), use points selected by the solver. + dense_output : bool, optional + Whether to compute a continuous solution. Default is False. + events : callable, or list of callables, optional + Events to track. If None (default), no events will be tracked. + Each event occurs at the zeros of a continuous function of time and + state. Each function must have the signature ``event(t, y)`` and return + a float. The solver will find an accurate value of `t` at which + ``event(t, y(t)) = 0`` using a root-finding algorithm. By default, all + zeros will be found. The solver looks for a sign change over each step, + so if multiple zero crossings occur within one step, events may be + missed. Additionally each `event` function might have the following + attributes: + + terminal: bool, optional + Whether to terminate integration if this event occurs. + Implicitly False if not assigned. + direction: float, optional + Direction of a zero crossing. If `direction` is positive, + `event` will only trigger when going from negative to positive, + and vice versa if `direction` is negative. If 0, then either + direction will trigger event. Implicitly 0 if not assigned. + + You can assign attributes like ``event.terminal = True`` to any + function in Python. + vectorized : bool, optional + Whether `fun` can be called in a vectorized fashion. Default is False. + + If ``vectorized`` is False, `fun` will always be called with ``y`` of + shape ``(n,)``, where ``n = len(y0)``. + + If ``vectorized`` is True, `fun` may be called with ``y`` of shape + ``(n, k)``, where ``k`` is an integer. In this case, `fun` must behave + such that ``fun(t, y)[:, i] == fun(t, y[:, i])`` (i.e. each column of + the returned array is the time derivative of the state corresponding + with a column of ``y``). + + Setting ``vectorized=True`` allows for faster finite difference + approximation of the Jacobian by methods 'Radau' and 'BDF', but + will result in slower execution for other methods and for 'Radau' and + 'BDF' in some circumstances (e.g. small ``len(y0)``). + args : tuple, optional + Additional arguments to pass to the user-defined functions. If given, + the additional arguments are passed to all user-defined functions. + So if, for example, `fun` has the signature ``fun(t, y, a, b, c)``, + then `jac` (if given) and any event functions must have the same + signature, and `args` must be a tuple of length 3. + **options + Options passed to a chosen solver. All options available for already + implemented solvers are listed below. + first_step : float or None, optional + Initial step size. Default is `None` which means that the algorithm + should choose. + max_step : float, optional + Maximum allowed step size. Default is np.inf, i.e., the step size is not + bounded and determined solely by the solver. + rtol, atol : float or array_like, optional + Relative and absolute tolerances. The solver keeps the local error + estimates less than ``atol + rtol * abs(y)``. Here `rtol` controls a + relative accuracy (number of correct digits), while `atol` controls + absolute accuracy (number of correct decimal places). To achieve the + desired `rtol`, set `atol` to be smaller than the smallest value that + can be expected from ``rtol * abs(y)`` so that `rtol` dominates the + allowable error. If `atol` is larger than ``rtol * abs(y)`` the + number of correct digits is not guaranteed. Conversely, to achieve the + desired `atol` set `rtol` such that ``rtol * abs(y)`` is always smaller + than `atol`. If components of y have different scales, it might be + beneficial to set different `atol` values for different components by + passing array_like with shape (n,) for `atol`. Default values are + 1e-3 for `rtol` and 1e-6 for `atol`. + jac : array_like, sparse_matrix, callable or None, optional + Jacobian matrix of the right-hand side of the system with respect + to y, required by the 'Radau', 'BDF' and 'LSODA' method. The + Jacobian matrix has shape (n, n) and its element (i, j) is equal to + ``d f_i / d y_j``. There are three ways to define the Jacobian: + + * If array_like or sparse_matrix, the Jacobian is assumed to + be constant. Not supported by 'LSODA'. + * If callable, the Jacobian is assumed to depend on both + t and y; it will be called as ``jac(t, y)``, as necessary. + For 'Radau' and 'BDF' methods, the return value might be a + sparse matrix. + * If None (default), the Jacobian will be approximated by + finite differences. + + It is generally recommended to provide the Jacobian rather than + relying on a finite-difference approximation. + jac_sparsity : array_like, sparse matrix or None, optional + Defines a sparsity structure of the Jacobian matrix for a finite- + difference approximation. Its shape must be (n, n). This argument + is ignored if `jac` is not `None`. If the Jacobian has only few + non-zero elements in *each* row, providing the sparsity structure + will greatly speed up the computations [10]_. A zero entry means that + a corresponding element in the Jacobian is always zero. If None + (default), the Jacobian is assumed to be dense. + Not supported by 'LSODA', see `lband` and `uband` instead. + lband, uband : int or None, optional + Parameters defining the bandwidth of the Jacobian for the 'LSODA' + method, i.e., ``jac[i, j] != 0 only for i - lband <= j <= i + uband``. + Default is None. Setting these requires your jac routine to return the + Jacobian in the packed format: the returned array must have ``n`` + columns and ``uband + lband + 1`` rows in which Jacobian diagonals are + written. Specifically ``jac_packed[uband + i - j , j] = jac[i, j]``. + The same format is used in `scipy.linalg.solve_banded` (check for an + illustration). These parameters can be also used with ``jac=None`` to + reduce the number of Jacobian elements estimated by finite differences. + min_step : float, optional + The minimum allowed step size for 'LSODA' method. + By default `min_step` is zero. + + Returns + ------- + Bunch object with the following fields defined: + t : ndarray, shape (n_points,) + Time points. + y : ndarray, shape (n, n_points) + Values of the solution at `t`. + sol : `OdeSolution` or None + Found solution as `OdeSolution` instance; None if `dense_output` was + set to False. + t_events : list of ndarray or None + Contains for each event type a list of arrays at which an event of + that type event was detected. None if `events` was None. + y_events : list of ndarray or None + For each value of `t_events`, the corresponding value of the solution. + None if `events` was None. + nfev : int + Number of evaluations of the right-hand side. + njev : int + Number of evaluations of the Jacobian. + nlu : int + Number of LU decompositions. + status : int + Reason for algorithm termination: + + * -1: Integration step failed. + * 0: The solver successfully reached the end of `tspan`. + * 1: A termination event occurred. + + message : string + Human-readable description of the termination reason. + success : bool + True if the solver reached the interval end or a termination event + occurred (``status >= 0``). + + References + ---------- + .. [1] J. R. Dormand, P. J. Prince, "A family of embedded Runge-Kutta + formulae", Journal of Computational and Applied Mathematics, Vol. 6, + No. 1, pp. 19-26, 1980. + .. [2] L. W. Shampine, "Some Practical Runge-Kutta Formulas", Mathematics + of Computation,, Vol. 46, No. 173, pp. 135-150, 1986. + .. [3] P. Bogacki, L.F. Shampine, "A 3(2) Pair of Runge-Kutta Formulas", + Appl. Math. Lett. Vol. 2, No. 4. pp. 321-325, 1989. + .. [4] E. Hairer, G. Wanner, "Solving Ordinary Differential Equations II: + Stiff and Differential-Algebraic Problems", Sec. IV.8. + .. [5] `Backward Differentiation Formula + `_ + on Wikipedia. + .. [6] L. F. Shampine, M. W. Reichelt, "THE MATLAB ODE SUITE", SIAM J. SCI. + COMPUTE., Vol. 18, No. 1, pp. 1-22, January 1997. + .. [7] A. C. Hindmarsh, "ODEPACK, A Systematized Collection of ODE + Solvers," IMACS Transactions on Scientific Computation, Vol 1., + pp. 55-64, 1983. + .. [8] L. Petzold, "Automatic selection of methods for solving stiff and + nonstiff systems of ordinary differential equations", SIAM Journal + on Scientific and Statistical Computing, Vol. 4, No. 1, pp. 136-148, + 1983. + .. [9] `Stiff equation `_ on + Wikipedia. + .. [10] A. Curtis, M. J. D. Powell, and J. Reid, "On the estimation of + sparse Jacobian matrices", Journal of the Institute of Mathematics + and its Applications, 13, pp. 117-120, 1974. + .. [11] `Cauchy-Riemann equations + `_ on + Wikipedia. + .. [12] `Lotka-Volterra equations + `_ + on Wikipedia. + .. [13] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential + Equations I: Nonstiff Problems", Sec. II. + .. [14] `Page with original Fortran code of DOP853 + `_. + + Examples + -------- + Basic exponential decay showing automatically chosen time points. + + >>> import numpy as np + >>> from scipy.integrate import solve_ivp + >>> def exponential_decay(t, y): return -0.5 * y + >>> sol = solve_ivp(exponential_decay, [0, 10], [2, 4, 8]) + >>> print(sol.t) + [ 0. 0.11487653 1.26364188 3.06061781 4.81611105 6.57445806 + 8.33328988 10. ] + >>> print(sol.y) + [[2. 1.88836035 1.06327177 0.43319312 0.18017253 0.07483045 + 0.03107158 0.01350781] + [4. 3.7767207 2.12654355 0.86638624 0.36034507 0.14966091 + 0.06214316 0.02701561] + [8. 7.5534414 4.25308709 1.73277247 0.72069014 0.29932181 + 0.12428631 0.05403123]] + + Specifying points where the solution is desired. + + >>> sol = solve_ivp(exponential_decay, [0, 10], [2, 4, 8], + ... t_eval=[0, 1, 2, 4, 10]) + >>> print(sol.t) + [ 0 1 2 4 10] + >>> print(sol.y) + [[2. 1.21305369 0.73534021 0.27066736 0.01350938] + [4. 2.42610739 1.47068043 0.54133472 0.02701876] + [8. 4.85221478 2.94136085 1.08266944 0.05403753]] + + Cannon fired upward with terminal event upon impact. The ``terminal`` and + ``direction`` fields of an event are applied by monkey patching a function. + Here ``y[0]`` is position and ``y[1]`` is velocity. The projectile starts + at position 0 with velocity +10. Note that the integration never reaches + t=100 because the event is terminal. + + >>> def upward_cannon(t, y): return [y[1], -0.5] + >>> def hit_ground(t, y): return y[0] + >>> hit_ground.terminal = True + >>> hit_ground.direction = -1 + >>> sol = solve_ivp(upward_cannon, [0, 100], [0, 10], events=hit_ground) + >>> print(sol.t_events) + [array([40.])] + >>> print(sol.t) + [0.00000000e+00 9.99900010e-05 1.09989001e-03 1.10988901e-02 + 1.11088891e-01 1.11098890e+00 1.11099890e+01 4.00000000e+01] + + Use `dense_output` and `events` to find position, which is 100, at the apex + of the cannonball's trajectory. Apex is not defined as terminal, so both + apex and hit_ground are found. There is no information at t=20, so the sol + attribute is used to evaluate the solution. The sol attribute is returned + by setting ``dense_output=True``. Alternatively, the `y_events` attribute + can be used to access the solution at the time of the event. + + >>> def apex(t, y): return y[1] + >>> sol = solve_ivp(upward_cannon, [0, 100], [0, 10], + ... events=(hit_ground, apex), dense_output=True) + >>> print(sol.t_events) + [array([40.]), array([20.])] + >>> print(sol.t) + [0.00000000e+00 9.99900010e-05 1.09989001e-03 1.10988901e-02 + 1.11088891e-01 1.11098890e+00 1.11099890e+01 4.00000000e+01] + >>> print(sol.sol(sol.t_events[1][0])) + [100. 0.] + >>> print(sol.y_events) + [array([[-5.68434189e-14, -1.00000000e+01]]), array([[1.00000000e+02, 1.77635684e-15]])] + + As an example of a system with additional parameters, we'll implement + the Lotka-Volterra equations [12]_. + + >>> def lotkavolterra(t, z, a, b, c, d): + ... x, y = z + ... return [a*x - b*x*y, -c*y + d*x*y] + ... + + We pass in the parameter values a=1.5, b=1, c=3 and d=1 with the `args` + argument. + + >>> sol = solve_ivp(lotkavolterra, [0, 15], [10, 5], args=(1.5, 1, 3, 1), + ... dense_output=True) + + Compute a dense solution and plot it. + + >>> t = np.linspace(0, 15, 300) + >>> z = sol.sol(t) + >>> import matplotlib.pyplot as plt + >>> plt.plot(t, z.T) + >>> plt.xlabel('t') + >>> plt.legend(['x', 'y'], shadow=True) + >>> plt.title('Lotka-Volterra System') + >>> plt.show() + """ + # if method not in METHODS and not ( + # inspect.isclass(method) and issubclass(method, OdeSolver)): + # raise ValueError("`method` must be one of {} or OdeSolver class." + # .format(METHODS)) + + t0, tf = map(float, t_span) + + if args is not None: + # Wrap the user's fun (and jac, if given) in lambdas to hide the + # additional parameters. Pass in the original fun as a keyword + # argument to keep it in the scope of the lambda. + try: + _ = [*(args)] + except TypeError as exp: + suggestion_tuple = ( + "Supplied 'args' cannot be unpacked. Please supply `args`" + f" as a tuple (e.g. `args=({args},)`)" + ) + raise TypeError(suggestion_tuple) from exp + + def fun(t, x, fun=fun): + return fun(t, x, *args) + + jac = options.get("jac") + if callable(jac): + options["jac"] = lambda t, x: jac(t, x, *args) + + if t_eval is not None: + t_eval = np.asarray(t_eval) + if t_eval.ndim != 1: + raise ValueError("`t_eval` must be 1-dimensional.") + + if np.any(t_eval < min(t0, tf)) or np.any(t_eval > max(t0, tf)): + raise ValueError("Values in `t_eval` are not within `t_span`.") + + d = np.diff(t_eval) + if tf > t0 and np.any(d <= 0) or tf < t0 and np.any(d >= 0): + raise ValueError("Values in `t_eval` are not properly sorted.") + + if tf > t0: + t_eval_i = 0 + else: + # Make order of t_eval decreasing to use np.searchsorted. + t_eval = t_eval[::-1] + # This will be an upper bound for slices. + t_eval_i = t_eval.shape[0] + + # if method in METHODS: + # method = METHODS[method] + + solver = method(fun, t0, y0, tf, vectorized=vectorized, **options) + + if t_eval is None: + ts = [t0] + ys = [y0] + elif t_eval is not None and dense_output: + ts = [] + ti = [t0] + ys = [] + else: + ts = [] + ys = [] + + interpolants = [] + + events, is_terminal, event_dir = prepare_events(events) + + if events is not None: + if args is not None: + # Wrap user functions in lambdas to hide the additional parameters. + # The original event function is passed as a keyword argument to the + # lambda to keep the original function in scope (i.e., avoid the + # late binding closure "gotcha"). + events = [lambda t, x, event=event: event(t, x, *args) for event in events] + g = [event(t0, y0) for event in events] + t_events = [[] for _ in range(len(events))] + y_events = [[] for _ in range(len(events))] + else: + t_events = None + y_events = None + + status = None + while status is None: + message = solver.step() + + if solver.status == "finished": + status = 0 + elif solver.status == "failed": + status = -1 + break + + t_old = solver.t_old + t = solver.t + y = solver.y + + if dense_output: + sol = solver.dense_output() + interpolants.append(sol) + else: + sol = None + + if events is not None: + g_new = [event(t, y) for event in events] + active_events = find_active_events(g, g_new, event_dir) + if active_events.size > 0: + if sol is None: + sol = solver.dense_output() + + root_indices, roots, terminate = handle_events( + sol, events, active_events, is_terminal, t_old, t + ) + + for e, te in zip(root_indices, roots): + t_events[e].append(te) + y_events[e].append(sol(te)) + + if terminate: + status = 1 + t = roots[-1] + y = sol(t) + + g = g_new + + if t_eval is None: + ts.append(t) + ys.append(y) + else: + # The value in t_eval equal to t will be included. + if solver.direction > 0: + t_eval_i_new = np.searchsorted(t_eval, t, side="right") + t_eval_step = t_eval[t_eval_i:t_eval_i_new] + else: + t_eval_i_new = np.searchsorted(t_eval, t, side="left") + # It has to be done with two slice operations, because + # you can't slice to 0th element inclusive using backward + # slicing. + t_eval_step = t_eval[t_eval_i_new:t_eval_i][::-1] + + if t_eval_step.size > 0: + if sol is None: + sol = solver.dense_output() + ts.append(t_eval_step) + ys.append(sol(t_eval_step)) + t_eval_i = t_eval_i_new + + if t_eval is not None and dense_output: + ti.append(t) + + message = MESSAGES.get(status, message) + + if t_events is not None: + t_events = [np.asarray(te) for te in t_events] + y_events = [np.asarray(ye) for ye in y_events] + + if t_eval is None: + ts = np.array(ts) + ys = np.vstack(ys).T + elif ts: + ts = np.hstack(ts) + ys = np.hstack(ys) + + if dense_output: + if t_eval is None: + sol = OdeSolution(ts, interpolants) + else: + sol = OdeSolution(ti, interpolants) + else: + sol = None - return solve_ivp_( - *args, - method=DOP853, - **kwargs, + return OdeResult( + t=ts, + y=ys, + sol=sol, + t_events=t_events, + y_events=y_events, + nfev=solver.nfev, + njev=solver.njev, + nlu=solver.nlu, + status=status, + message=message, + success=status >= 0, ) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py new file mode 100644 index 000000000..e69de29bb From 0fa2a4949a1318751a241e03dbc8d2f53129151f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 9 Jan 2024 13:57:07 +0100 Subject: [PATCH 079/346] more ivp draft --- src/hapsira/core/math/ivp/_api.py | 198 +------ src/hapsira/core/math/ivp/_base.py | 291 ++++++++++ src/hapsira/core/math/ivp/_brentq.py | 300 +++++++++++ src/hapsira/core/math/ivp/_common.py | 122 +++++ .../core/math/ivp/_dop853_coefficients.py | 193 +++++++ src/hapsira/core/math/ivp/_rk.py | 496 ++++++++++++++++++ 6 files changed, 1411 insertions(+), 189 deletions(-) create mode 100644 src/hapsira/core/math/ivp/_base.py create mode 100644 src/hapsira/core/math/ivp/_brentq.py create mode 100644 src/hapsira/core/math/ivp/_common.py create mode 100644 src/hapsira/core/math/ivp/_dop853_coefficients.py diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 5dc35a123..e7a74627a 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -1,7 +1,15 @@ -from scipy.integrate import DOP853 import numpy as np +from ._brentq import brentq +from ._common import OdeSolution + +from ._rk import EPS, DOP853 + + +MESSAGES = {0: "The solver successfully reached the end of the integration interval.", + 1: "A termination event occurred."} + class OdeResult(dict): """Represents the optimization result. @@ -108,7 +116,6 @@ def solve_event_equation(event, sol, t_old, t): root : float Found solution. """ - from scipy.optimize import brentq return brentq(lambda t: event(t, sol(t)), t_old, t, xtol=4 * EPS, rtol=4 * EPS) # https://github.com/scipy/scipy/blob/ea4d1f1330950bee74396e427fe6330424907621/scipy/optimize/_zeros_py.py#L682 @@ -229,25 +236,6 @@ def solve_ivp( ): """Solve an initial value problem for a system of ODEs. - This function numerically integrates a system of ordinary differential - equations given an initial value:: - - dy / dt = f(t, y) - y(t0) = y0 - - Here t is a 1-D independent variable (time), y(t) is an - N-D vector-valued function (state), and an N-D - vector-valued function f(t, y) determines the differential equations. - The goal is to find y(t) approximately satisfying the differential - equations, given an initial value y(t0)=y0. - - Some of the solvers support integration in the complex domain, but note - that for stiff ODE solvers, the right-hand side must be - complex-differentiable (satisfy Cauchy-Riemann equations [11]_). - To solve a problem in the complex domain, pass y0 with a complex data type. - Another option always available is to rewrite your problem for real and - imaginary parts separately. - Parameters ---------- fun : callable @@ -265,48 +253,11 @@ def solve_ivp( complex data type (even if the initial value is purely real). method : string or `OdeSolver`, optional Integration method to use: - - * 'RK45' (default): Explicit Runge-Kutta method of order 5(4) [1]_. - The error is controlled assuming accuracy of the fourth-order - method, but steps are taken using the fifth-order accurate - formula (local extrapolation is done). A quartic interpolation - polynomial is used for the dense output [2]_. Can be applied in - the complex domain. - * 'RK23': Explicit Runge-Kutta method of order 3(2) [3]_. The error - is controlled assuming accuracy of the second-order method, but - steps are taken using the third-order accurate formula (local - extrapolation is done). A cubic Hermite polynomial is used for the - dense output. Can be applied in the complex domain. * 'DOP853': Explicit Runge-Kutta method of order 8 [13]_. Python implementation of the "DOP853" algorithm originally written in Fortran [14]_. A 7-th order interpolation polynomial accurate to 7-th order is used for the dense output. Can be applied in the complex domain. - * 'Radau': Implicit Runge-Kutta method of the Radau IIA family of - order 5 [4]_. The error is controlled with a third-order accurate - embedded formula. A cubic polynomial which satisfies the - collocation conditions is used for the dense output. - * 'BDF': Implicit multi-step variable-order (1 to 5) method based - on a backward differentiation formula for the derivative - approximation [5]_. The implementation follows the one described - in [6]_. A quasi-constant step scheme is used and accuracy is - enhanced using the NDF modification. Can be applied in the - complex domain. - * 'LSODA': Adams/BDF method with automatic stiffness detection and - switching [7]_, [8]_. This is a wrapper of the Fortran solver - from ODEPACK. - - Explicit Runge-Kutta methods ('RK23', 'RK45', 'DOP853') should be used - for non-stiff problems and implicit methods ('Radau', 'BDF') for - stiff problems [9]_. Among Runge-Kutta methods, 'DOP853' is recommended - for solving with high precision (low values of `rtol` and `atol`). - - If not sure, first try to run 'RK45'. If it makes unusually many - iterations, diverges, or fails, your problem is likely to be stiff and - you should use 'Radau' or 'BDF'. 'LSODA' can also be a good universal - choice, but it might be somewhat less convenient to work with as it - wraps old Fortran code. - You can also pass an arbitrary class derived from `OdeSolver` which implements the solver. t_eval : array_like or None, optional @@ -456,137 +407,6 @@ def solve_ivp( True if the solver reached the interval end or a termination event occurred (``status >= 0``). - References - ---------- - .. [1] J. R. Dormand, P. J. Prince, "A family of embedded Runge-Kutta - formulae", Journal of Computational and Applied Mathematics, Vol. 6, - No. 1, pp. 19-26, 1980. - .. [2] L. W. Shampine, "Some Practical Runge-Kutta Formulas", Mathematics - of Computation,, Vol. 46, No. 173, pp. 135-150, 1986. - .. [3] P. Bogacki, L.F. Shampine, "A 3(2) Pair of Runge-Kutta Formulas", - Appl. Math. Lett. Vol. 2, No. 4. pp. 321-325, 1989. - .. [4] E. Hairer, G. Wanner, "Solving Ordinary Differential Equations II: - Stiff and Differential-Algebraic Problems", Sec. IV.8. - .. [5] `Backward Differentiation Formula - `_ - on Wikipedia. - .. [6] L. F. Shampine, M. W. Reichelt, "THE MATLAB ODE SUITE", SIAM J. SCI. - COMPUTE., Vol. 18, No. 1, pp. 1-22, January 1997. - .. [7] A. C. Hindmarsh, "ODEPACK, A Systematized Collection of ODE - Solvers," IMACS Transactions on Scientific Computation, Vol 1., - pp. 55-64, 1983. - .. [8] L. Petzold, "Automatic selection of methods for solving stiff and - nonstiff systems of ordinary differential equations", SIAM Journal - on Scientific and Statistical Computing, Vol. 4, No. 1, pp. 136-148, - 1983. - .. [9] `Stiff equation `_ on - Wikipedia. - .. [10] A. Curtis, M. J. D. Powell, and J. Reid, "On the estimation of - sparse Jacobian matrices", Journal of the Institute of Mathematics - and its Applications, 13, pp. 117-120, 1974. - .. [11] `Cauchy-Riemann equations - `_ on - Wikipedia. - .. [12] `Lotka-Volterra equations - `_ - on Wikipedia. - .. [13] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential - Equations I: Nonstiff Problems", Sec. II. - .. [14] `Page with original Fortran code of DOP853 - `_. - - Examples - -------- - Basic exponential decay showing automatically chosen time points. - - >>> import numpy as np - >>> from scipy.integrate import solve_ivp - >>> def exponential_decay(t, y): return -0.5 * y - >>> sol = solve_ivp(exponential_decay, [0, 10], [2, 4, 8]) - >>> print(sol.t) - [ 0. 0.11487653 1.26364188 3.06061781 4.81611105 6.57445806 - 8.33328988 10. ] - >>> print(sol.y) - [[2. 1.88836035 1.06327177 0.43319312 0.18017253 0.07483045 - 0.03107158 0.01350781] - [4. 3.7767207 2.12654355 0.86638624 0.36034507 0.14966091 - 0.06214316 0.02701561] - [8. 7.5534414 4.25308709 1.73277247 0.72069014 0.29932181 - 0.12428631 0.05403123]] - - Specifying points where the solution is desired. - - >>> sol = solve_ivp(exponential_decay, [0, 10], [2, 4, 8], - ... t_eval=[0, 1, 2, 4, 10]) - >>> print(sol.t) - [ 0 1 2 4 10] - >>> print(sol.y) - [[2. 1.21305369 0.73534021 0.27066736 0.01350938] - [4. 2.42610739 1.47068043 0.54133472 0.02701876] - [8. 4.85221478 2.94136085 1.08266944 0.05403753]] - - Cannon fired upward with terminal event upon impact. The ``terminal`` and - ``direction`` fields of an event are applied by monkey patching a function. - Here ``y[0]`` is position and ``y[1]`` is velocity. The projectile starts - at position 0 with velocity +10. Note that the integration never reaches - t=100 because the event is terminal. - - >>> def upward_cannon(t, y): return [y[1], -0.5] - >>> def hit_ground(t, y): return y[0] - >>> hit_ground.terminal = True - >>> hit_ground.direction = -1 - >>> sol = solve_ivp(upward_cannon, [0, 100], [0, 10], events=hit_ground) - >>> print(sol.t_events) - [array([40.])] - >>> print(sol.t) - [0.00000000e+00 9.99900010e-05 1.09989001e-03 1.10988901e-02 - 1.11088891e-01 1.11098890e+00 1.11099890e+01 4.00000000e+01] - - Use `dense_output` and `events` to find position, which is 100, at the apex - of the cannonball's trajectory. Apex is not defined as terminal, so both - apex and hit_ground are found. There is no information at t=20, so the sol - attribute is used to evaluate the solution. The sol attribute is returned - by setting ``dense_output=True``. Alternatively, the `y_events` attribute - can be used to access the solution at the time of the event. - - >>> def apex(t, y): return y[1] - >>> sol = solve_ivp(upward_cannon, [0, 100], [0, 10], - ... events=(hit_ground, apex), dense_output=True) - >>> print(sol.t_events) - [array([40.]), array([20.])] - >>> print(sol.t) - [0.00000000e+00 9.99900010e-05 1.09989001e-03 1.10988901e-02 - 1.11088891e-01 1.11098890e+00 1.11099890e+01 4.00000000e+01] - >>> print(sol.sol(sol.t_events[1][0])) - [100. 0.] - >>> print(sol.y_events) - [array([[-5.68434189e-14, -1.00000000e+01]]), array([[1.00000000e+02, 1.77635684e-15]])] - - As an example of a system with additional parameters, we'll implement - the Lotka-Volterra equations [12]_. - - >>> def lotkavolterra(t, z, a, b, c, d): - ... x, y = z - ... return [a*x - b*x*y, -c*y + d*x*y] - ... - - We pass in the parameter values a=1.5, b=1, c=3 and d=1 with the `args` - argument. - - >>> sol = solve_ivp(lotkavolterra, [0, 15], [10, 5], args=(1.5, 1, 3, 1), - ... dense_output=True) - - Compute a dense solution and plot it. - - >>> t = np.linspace(0, 15, 300) - >>> z = sol.sol(t) - >>> import matplotlib.pyplot as plt - >>> plt.plot(t, z.T) - >>> plt.xlabel('t') - >>> plt.legend(['x', 'y'], shadow=True) - >>> plt.title('Lotka-Volterra System') - >>> plt.show() - """ # if method not in METHODS and not ( # inspect.isclass(method) and issubclass(method, OdeSolver)): diff --git a/src/hapsira/core/math/ivp/_base.py b/src/hapsira/core/math/ivp/_base.py new file mode 100644 index 000000000..05c00e60c --- /dev/null +++ b/src/hapsira/core/math/ivp/_base.py @@ -0,0 +1,291 @@ + +import numpy as np + + +def check_arguments(fun, y0, support_complex): + """Helper function for checking arguments common to all solvers.""" + y0 = np.asarray(y0) + if np.issubdtype(y0.dtype, np.complexfloating): + if not support_complex: + raise ValueError("`y0` is complex, but the chosen solver does " + "not support integration in a complex domain.") + dtype = complex + else: + dtype = float + y0 = y0.astype(dtype, copy=False) + + if y0.ndim != 1: + raise ValueError("`y0` must be 1-dimensional.") + + if not np.isfinite(y0).all(): + raise ValueError("All components of the initial state `y0` must be finite.") + + def fun_wrapped(t, y): + return np.asarray(fun(t, y), dtype=dtype) + + return fun_wrapped, y0 + + +class OdeSolver: + """Base class for ODE solvers. + + In order to implement a new solver you need to follow the guidelines: + + 1. A constructor must accept parameters presented in the base class + (listed below) along with any other parameters specific to a solver. + 2. A constructor must accept arbitrary extraneous arguments + ``**extraneous``, but warn that these arguments are irrelevant + using `common.warn_extraneous` function. Do not pass these + arguments to the base class. + 3. A solver must implement a private method `_step_impl(self)` which + propagates a solver one step further. It must return tuple + ``(success, message)``, where ``success`` is a boolean indicating + whether a step was successful, and ``message`` is a string + containing description of a failure if a step failed or None + otherwise. + 4. A solver must implement a private method `_dense_output_impl(self)`, + which returns a `DenseOutput` object covering the last successful + step. + 5. A solver must have attributes listed below in Attributes section. + Note that ``t_old`` and ``step_size`` are updated automatically. + 6. Use `fun(self, t, y)` method for the system rhs evaluation, this + way the number of function evaluations (`nfev`) will be tracked + automatically. + 7. For convenience, a base class provides `fun_single(self, t, y)` and + `fun_vectorized(self, t, y)` for evaluating the rhs in + non-vectorized and vectorized fashions respectively (regardless of + how `fun` from the constructor is implemented). These calls don't + increment `nfev`. + 8. If a solver uses a Jacobian matrix and LU decompositions, it should + track the number of Jacobian evaluations (`njev`) and the number of + LU decompositions (`nlu`). + 9. By convention, the function evaluations used to compute a finite + difference approximation of the Jacobian should not be counted in + `nfev`, thus use `fun_single(self, t, y)` or + `fun_vectorized(self, t, y)` when computing a finite difference + approximation of the Jacobian. + + Parameters + ---------- + fun : callable + Right-hand side of the system: the time derivative of the state ``y`` + at time ``t``. The calling signature is ``fun(t, y)``, where ``t`` is a + scalar and ``y`` is an ndarray with ``len(y) = len(y0)``. ``fun`` must + return an array of the same shape as ``y``. See `vectorized` for more + information. + t0 : float + Initial time. + y0 : array_like, shape (n,) + Initial state. + t_bound : float + Boundary time --- the integration won't continue beyond it. It also + determines the direction of the integration. + vectorized : bool + Whether `fun` can be called in a vectorized fashion. Default is False. + + If ``vectorized`` is False, `fun` will always be called with ``y`` of + shape ``(n,)``, where ``n = len(y0)``. + + If ``vectorized`` is True, `fun` may be called with ``y`` of shape + ``(n, k)``, where ``k`` is an integer. In this case, `fun` must behave + such that ``fun(t, y)[:, i] == fun(t, y[:, i])`` (i.e. each column of + the returned array is the time derivative of the state corresponding + with a column of ``y``). + + Setting ``vectorized=True`` allows for faster finite difference + approximation of the Jacobian by methods 'Radau' and 'BDF', but + will result in slower execution for other methods. It can also + result in slower overall execution for 'Radau' and 'BDF' in some + circumstances (e.g. small ``len(y0)``). + support_complex : bool, optional + Whether integration in a complex domain should be supported. + Generally determined by a derived solver class capabilities. + Default is False. + + Attributes + ---------- + n : int + Number of equations. + status : string + Current status of the solver: 'running', 'finished' or 'failed'. + t_bound : float + Boundary time. + direction : float + Integration direction: +1 or -1. + t : float + Current time. + y : ndarray + Current state. + t_old : float + Previous time. None if no steps were made yet. + step_size : float + Size of the last successful step. None if no steps were made yet. + nfev : int + Number of the system's rhs evaluations. + njev : int + Number of the Jacobian evaluations. + nlu : int + Number of LU decompositions. + """ + TOO_SMALL_STEP = "Required step size is less than spacing between numbers." + + def __init__(self, fun, t0, y0, t_bound, vectorized, + support_complex=False): + self.t_old = None + self.t = t0 + self._fun, self.y = check_arguments(fun, y0, support_complex) + self.t_bound = t_bound + self.vectorized = vectorized + + if vectorized: + def fun_single(t, y): + return self._fun(t, y[:, None]).ravel() + fun_vectorized = self._fun + else: + fun_single = self._fun + + def fun_vectorized(t, y): + f = np.empty_like(y) + for i, yi in enumerate(y.T): + f[:, i] = self._fun(t, yi) + return f + + def fun(t, y): + self.nfev += 1 + return self.fun_single(t, y) + + self.fun = fun + self.fun_single = fun_single + self.fun_vectorized = fun_vectorized + + self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 + self.n = self.y.size + self.status = 'running' + + self.nfev = 0 + self.njev = 0 + self.nlu = 0 + + @property + def step_size(self): + if self.t_old is None: + return None + else: + return np.abs(self.t - self.t_old) + + def step(self): + """Perform one integration step. + + Returns + ------- + message : string or None + Report from the solver. Typically a reason for a failure if + `self.status` is 'failed' after the step was taken or None + otherwise. + """ + if self.status != 'running': + raise RuntimeError("Attempt to step on a failed or finished " + "solver.") + + if self.n == 0 or self.t == self.t_bound: + # Handle corner cases of empty solver or no integration. + self.t_old = self.t + self.t = self.t_bound + message = None + self.status = 'finished' + else: + t = self.t + success, message = self._step_impl() + + if not success: + self.status = 'failed' + else: + self.t_old = t + if self.direction * (self.t - self.t_bound) >= 0: + self.status = 'finished' + + return message + + def dense_output(self): + """Compute a local interpolant over the last successful step. + + Returns + ------- + sol : `DenseOutput` + Local interpolant over the last successful step. + """ + if self.t_old is None: + raise RuntimeError("Dense output is available after a successful " + "step was made.") + + if self.n == 0 or self.t == self.t_old: + # Handle corner cases of empty solver and no integration. + return ConstantDenseOutput(self.t_old, self.t, self.y) + else: + return self._dense_output_impl() + + def _step_impl(self): + raise NotImplementedError + + def _dense_output_impl(self): + raise NotImplementedError + + +class DenseOutput: + """Base class for local interpolant over step made by an ODE solver. + + It interpolates between `t_min` and `t_max` (see Attributes below). + Evaluation outside this interval is not forbidden, but the accuracy is not + guaranteed. + + Attributes + ---------- + t_min, t_max : float + Time range of the interpolation. + """ + def __init__(self, t_old, t): + self.t_old = t_old + self.t = t + self.t_min = min(t, t_old) + self.t_max = max(t, t_old) + + def __call__(self, t): + """Evaluate the interpolant. + + Parameters + ---------- + t : float or array_like with shape (n_points,) + Points to evaluate the solution at. + + Returns + ------- + y : ndarray, shape (n,) or (n, n_points) + Computed values. Shape depends on whether `t` was a scalar or a + 1-D array. + """ + t = np.asarray(t) + if t.ndim > 1: + raise ValueError("`t` must be a float or a 1-D array.") + return self._call_impl(t) + + def _call_impl(self, t): + raise NotImplementedError + + +class ConstantDenseOutput(DenseOutput): + """Constant value interpolator. + + This class used for degenerate integration cases: equal integration limits + or a system with 0 equations. + """ + def __init__(self, t_old, t, value): + super().__init__(t_old, t) + self.value = value + + def _call_impl(self, t): + if t.ndim == 0: + return self.value + else: + ret = np.empty((self.value.shape[0], t.shape[0])) + ret[:] = self.value[:, None] + return ret diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py new file mode 100644 index 000000000..e7e072b6e --- /dev/null +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -0,0 +1,300 @@ + +import operator + +import numpy as np + + +# _zeros, OptimizeResult + + +_ECONVERGED = 0 +_ESIGNERR = -1 +_ECONVERR = -2 +_EVALUEERR = -3 +_EINPROGRESS = 1 + +CONVERGED = 'converged' +SIGNERR = 'sign error' +CONVERR = 'convergence error' +VALUEERR = 'value error' +INPROGRESS = 'No error' + +flag_map = {_ECONVERGED: CONVERGED, _ESIGNERR: SIGNERR, _ECONVERR: CONVERR, + _EVALUEERR: VALUEERR, _EINPROGRESS: INPROGRESS} + + +class OptimizeResult(dict): + """ Represents the optimization result. + + Attributes + ---------- + x : ndarray + The solution of the optimization. + success : bool + Whether or not the optimizer exited successfully. + status : int + Termination status of the optimizer. Its value depends on the + underlying solver. Refer to `message` for details. + message : str + Description of the cause of the termination. + fun, jac, hess: ndarray + Values of objective function, its Jacobian and its Hessian (if + available). The Hessians may be approximations, see the documentation + of the function in question. + hess_inv : object + Inverse of the objective function's Hessian; may be an approximation. + Not available for all solvers. The type of this attribute may be + either np.ndarray or scipy.sparse.linalg.LinearOperator. + nfev, njev, nhev : int + Number of evaluations of the objective functions and of its + Jacobian and Hessian. + nit : int + Number of iterations performed by the optimizer. + maxcv : float + The maximum constraint violation. + + Notes + ----- + Depending on the specific solver being used, `OptimizeResult` may + not have all attributes listed here, and they may have additional + attributes not listed here. Since this class is essentially a + subclass of dict with attribute accessors, one can see which + attributes are available using the `OptimizeResult.keys` method. + """ + + def __getattr__(self, name): + try: + return self[name] + except KeyError as e: + raise AttributeError(name) from e + + __setattr__ = dict.__setitem__ + __delattr__ = dict.__delitem__ + + # def __repr__(self): + # order_keys = ['message', 'success', 'status', 'fun', 'funl', 'x', 'xl', + # 'col_ind', 'nit', 'lower', 'upper', 'eqlin', 'ineqlin', + # 'converged', 'flag', 'function_calls', 'iterations', + # 'root'] + # # 'slack', 'con' are redundant with residuals + # # 'crossover_nit' is probably not interesting to most users + # omit_keys = {'slack', 'con', 'crossover_nit'} + + # def key(item): + # try: + # return order_keys.index(item[0].lower()) + # except ValueError: # item not in list + # return np.inf + + # def omit_redundant(items): + # for item in items: + # if item[0] in omit_keys: + # continue + # yield item + + # def item_sorter(d): + # return sorted(omit_redundant(d.items()), key=key) + + # if self.keys(): + # return _dict_formatter(self, sorter=item_sorter) + # else: + # return self.__class__.__name__ + "()" + + def __dir__(self): + return list(self.keys()) + + +class RootResults(OptimizeResult): + """Represents the root finding result. + + Attributes + ---------- + root : float + Estimated root location. + iterations : int + Number of iterations needed to find the root. + function_calls : int + Number of times the function was called. + converged : bool + True if the routine converged. + flag : str + Description of the cause of termination. + + """ + + def __init__(self, root, iterations, function_calls, flag): + self.root = root + self.iterations = iterations + self.function_calls = function_calls + self.converged = flag == _ECONVERGED + if flag in flag_map: + self.flag = flag_map[flag] + else: + self.flag = flag + + +def _wrap_nan_raise(f): + + def f_raise(x, *args): + fx = f(x, *args) + f_raise._function_calls += 1 + if np.isnan(fx): + msg = (f'The function value at x={x} is NaN; ' + 'solver cannot continue.') + err = ValueError(msg) + err._x = x + err._function_calls = f_raise._function_calls + raise err + return fx + + f_raise._function_calls = 0 + return f_raise + + +def results_c(full_output, r): + if full_output: + x, funcalls, iterations, flag = r + results = RootResults(root=x, + iterations=iterations, + function_calls=funcalls, + flag=flag) + return x, results + else: + return r + + +_iter = 100 +_xtol = 2e-12 +_rtol = 4 * np.finfo(float).eps + + +def brentq(f, a, b, args=(), + xtol=_xtol, rtol=_rtol, maxiter=_iter, + full_output=False, disp=True): + """ + Find a root of a function in a bracketing interval using Brent's method. + + Uses the classic Brent's method to find a root of the function `f` on + the sign changing interval [a , b]. Generally considered the best of the + rootfinding routines here. It is a safe version of the secant method that + uses inverse quadratic extrapolation. Brent's method combines root + bracketing, interval bisection, and inverse quadratic interpolation. It is + sometimes known as the van Wijngaarden-Dekker-Brent method. Brent (1973) + claims convergence is guaranteed for functions computable within [a,b]. + + [Brent1973]_ provides the classic description of the algorithm. Another + description can be found in a recent edition of Numerical Recipes, including + [PressEtal1992]_. A third description is at + http://mathworld.wolfram.com/BrentsMethod.html. It should be easy to + understand the algorithm just by reading our code. Our code diverges a bit + from standard presentations: we choose a different formula for the + extrapolation step. + + Parameters + ---------- + f : function + Python function returning a number. The function :math:`f` + must be continuous, and :math:`f(a)` and :math:`f(b)` must + have opposite signs. + a : scalar + One end of the bracketing interval :math:`[a, b]`. + b : scalar + The other end of the bracketing interval :math:`[a, b]`. + xtol : number, optional + The computed root ``x0`` will satisfy ``np.allclose(x, x0, + atol=xtol, rtol=rtol)``, where ``x`` is the exact root. The + parameter must be positive. For nice functions, Brent's + method will often satisfy the above condition with ``xtol/2`` + and ``rtol/2``. [Brent1973]_ + rtol : number, optional + The computed root ``x0`` will satisfy ``np.allclose(x, x0, + atol=xtol, rtol=rtol)``, where ``x`` is the exact root. The + parameter cannot be smaller than its default value of + ``4*np.finfo(float).eps``. For nice functions, Brent's + method will often satisfy the above condition with ``xtol/2`` + and ``rtol/2``. [Brent1973]_ + maxiter : int, optional + If convergence is not achieved in `maxiter` iterations, an error is + raised. Must be >= 0. + args : tuple, optional + Containing extra arguments for the function `f`. + `f` is called by ``apply(f, (x)+args)``. + full_output : bool, optional + If `full_output` is False, the root is returned. If `full_output` is + True, the return value is ``(x, r)``, where `x` is the root, and `r` is + a `RootResults` object. + disp : bool, optional + If True, raise RuntimeError if the algorithm didn't converge. + Otherwise, the convergence status is recorded in any `RootResults` + return object. + + Returns + ------- + root : float + Root of `f` between `a` and `b`. + r : `RootResults` (present if ``full_output = True``) + Object containing information about the convergence. In particular, + ``r.converged`` is True if the routine converged. + + Notes + ----- + `f` must be continuous. f(a) and f(b) must have opposite signs. + + Related functions fall into several classes: + + multivariate local optimizers + `fmin`, `fmin_powell`, `fmin_cg`, `fmin_bfgs`, `fmin_ncg` + nonlinear least squares minimizer + `leastsq` + constrained multivariate optimizers + `fmin_l_bfgs_b`, `fmin_tnc`, `fmin_cobyla` + global optimizers + `basinhopping`, `brute`, `differential_evolution` + local scalar minimizers + `fminbound`, `brent`, `golden`, `bracket` + N-D root-finding + `fsolve` + 1-D root-finding + `brenth`, `ridder`, `bisect`, `newton` + scalar fixed-point finder + `fixed_point` + + References + ---------- + .. [Brent1973] + Brent, R. P., + *Algorithms for Minimization Without Derivatives*. + Englewood Cliffs, NJ: Prentice-Hall, 1973. Ch. 3-4. + + .. [PressEtal1992] + Press, W. H.; Flannery, B. P.; Teukolsky, S. A.; and Vetterling, W. T. + *Numerical Recipes in FORTRAN: The Art of Scientific Computing*, 2nd ed. + Cambridge, England: Cambridge University Press, pp. 352-355, 1992. + Section 9.3: "Van Wijngaarden-Dekker-Brent Method." + + Examples + -------- + >>> def f(x): + ... return (x**2 - 1) + + >>> from scipy import optimize + + >>> root = optimize.brentq(f, -2, 0) + >>> root + -1.0 + + >>> root = optimize.brentq(f, 0, 2) + >>> root + 1.0 + """ + if not isinstance(args, tuple): + args = (args,) + maxiter = operator.index(maxiter) + if xtol <= 0: + raise ValueError("xtol too small (%g <= 0)" % xtol) + if rtol < _rtol: + raise ValueError(f"rtol too small ({rtol:g} < {_rtol:g})") + f = _wrap_nan_raise(f) + r = _zeros._brentq(f, a, b, xtol, rtol, maxiter, args, full_output, disp) + return results_c(full_output, r) + diff --git a/src/hapsira/core/math/ivp/_common.py b/src/hapsira/core/math/ivp/_common.py new file mode 100644 index 000000000..ad3bf5a45 --- /dev/null +++ b/src/hapsira/core/math/ivp/_common.py @@ -0,0 +1,122 @@ + +from itertools import groupby + +import numpy as np + + +class OdeSolution: + """Continuous ODE solution. + + It is organized as a collection of `DenseOutput` objects which represent + local interpolants. It provides an algorithm to select a right interpolant + for each given point. + + The interpolants cover the range between `t_min` and `t_max` (see + Attributes below). Evaluation outside this interval is not forbidden, but + the accuracy is not guaranteed. + + When evaluating at a breakpoint (one of the values in `ts`) a segment with + the lower index is selected. + + Parameters + ---------- + ts : array_like, shape (n_segments + 1,) + Time instants between which local interpolants are defined. Must + be strictly increasing or decreasing (zero segment with two points is + also allowed). + interpolants : list of DenseOutput with n_segments elements + Local interpolants. An i-th interpolant is assumed to be defined + between ``ts[i]`` and ``ts[i + 1]``. + + Attributes + ---------- + t_min, t_max : float + Time range of the interpolation. + """ + def __init__(self, ts, interpolants): + ts = np.asarray(ts) + d = np.diff(ts) + # The first case covers integration on zero segment. + if not ((ts.size == 2 and ts[0] == ts[-1]) + or np.all(d > 0) or np.all(d < 0)): + raise ValueError("`ts` must be strictly increasing or decreasing.") + + self.n_segments = len(interpolants) + if ts.shape != (self.n_segments + 1,): + raise ValueError("Numbers of time stamps and interpolants " + "don't match.") + + self.ts = ts + self.interpolants = interpolants + if ts[-1] >= ts[0]: + self.t_min = ts[0] + self.t_max = ts[-1] + self.ascending = True + self.ts_sorted = ts + else: + self.t_min = ts[-1] + self.t_max = ts[0] + self.ascending = False + self.ts_sorted = ts[::-1] + + def _call_single(self, t): + # Here we preserve a certain symmetry that when t is in self.ts, + # then we prioritize a segment with a lower index. + if self.ascending: + ind = np.searchsorted(self.ts_sorted, t, side='left') + else: + ind = np.searchsorted(self.ts_sorted, t, side='right') + + segment = min(max(ind - 1, 0), self.n_segments - 1) + if not self.ascending: + segment = self.n_segments - 1 - segment + + return self.interpolants[segment](t) + + def __call__(self, t): + """Evaluate the solution. + + Parameters + ---------- + t : float or array_like with shape (n_points,) + Points to evaluate at. + + Returns + ------- + y : ndarray, shape (n_states,) or (n_states, n_points) + Computed values. Shape depends on whether `t` is a scalar or a + 1-D array. + """ + t = np.asarray(t) + + if t.ndim == 0: + return self._call_single(t) + + order = np.argsort(t) + reverse = np.empty_like(order) + reverse[order] = np.arange(order.shape[0]) + t_sorted = t[order] + + # See comment in self._call_single. + if self.ascending: + segments = np.searchsorted(self.ts_sorted, t_sorted, side='left') + else: + segments = np.searchsorted(self.ts_sorted, t_sorted, side='right') + segments -= 1 + segments[segments < 0] = 0 + segments[segments > self.n_segments - 1] = self.n_segments - 1 + if not self.ascending: + segments = self.n_segments - 1 - segments + + ys = [] + group_start = 0 + for segment, group in groupby(segments): + group_end = group_start + len(list(group)) + y = self.interpolants[segment](t_sorted[group_start:group_end]) + ys.append(y) + group_start = group_end + + ys = np.hstack(ys) + ys = ys[:, reverse] + + return ys diff --git a/src/hapsira/core/math/ivp/_dop853_coefficients.py b/src/hapsira/core/math/ivp/_dop853_coefficients.py new file mode 100644 index 000000000..f39f2f365 --- /dev/null +++ b/src/hapsira/core/math/ivp/_dop853_coefficients.py @@ -0,0 +1,193 @@ +import numpy as np + +N_STAGES = 12 +N_STAGES_EXTENDED = 16 +INTERPOLATOR_POWER = 7 + +C = np.array([0.0, + 0.526001519587677318785587544488e-01, + 0.789002279381515978178381316732e-01, + 0.118350341907227396726757197510, + 0.281649658092772603273242802490, + 0.333333333333333333333333333333, + 0.25, + 0.307692307692307692307692307692, + 0.651282051282051282051282051282, + 0.6, + 0.857142857142857142857142857142, + 1.0, + 1.0, + 0.1, + 0.2, + 0.777777777777777777777777777778]) + +A = np.zeros((N_STAGES_EXTENDED, N_STAGES_EXTENDED)) +A[1, 0] = 5.26001519587677318785587544488e-2 + +A[2, 0] = 1.97250569845378994544595329183e-2 +A[2, 1] = 5.91751709536136983633785987549e-2 + +A[3, 0] = 2.95875854768068491816892993775e-2 +A[3, 2] = 8.87627564304205475450678981324e-2 + +A[4, 0] = 2.41365134159266685502369798665e-1 +A[4, 2] = -8.84549479328286085344864962717e-1 +A[4, 3] = 9.24834003261792003115737966543e-1 + +A[5, 0] = 3.7037037037037037037037037037e-2 +A[5, 3] = 1.70828608729473871279604482173e-1 +A[5, 4] = 1.25467687566822425016691814123e-1 + +A[6, 0] = 3.7109375e-2 +A[6, 3] = 1.70252211019544039314978060272e-1 +A[6, 4] = 6.02165389804559606850219397283e-2 +A[6, 5] = -1.7578125e-2 + +A[7, 0] = 3.70920001185047927108779319836e-2 +A[7, 3] = 1.70383925712239993810214054705e-1 +A[7, 4] = 1.07262030446373284651809199168e-1 +A[7, 5] = -1.53194377486244017527936158236e-2 +A[7, 6] = 8.27378916381402288758473766002e-3 + +A[8, 0] = 6.24110958716075717114429577812e-1 +A[8, 3] = -3.36089262944694129406857109825 +A[8, 4] = -8.68219346841726006818189891453e-1 +A[8, 5] = 2.75920996994467083049415600797e1 +A[8, 6] = 2.01540675504778934086186788979e1 +A[8, 7] = -4.34898841810699588477366255144e1 + +A[9, 0] = 4.77662536438264365890433908527e-1 +A[9, 3] = -2.48811461997166764192642586468 +A[9, 4] = -5.90290826836842996371446475743e-1 +A[9, 5] = 2.12300514481811942347288949897e1 +A[9, 6] = 1.52792336328824235832596922938e1 +A[9, 7] = -3.32882109689848629194453265587e1 +A[9, 8] = -2.03312017085086261358222928593e-2 + +A[10, 0] = -9.3714243008598732571704021658e-1 +A[10, 3] = 5.18637242884406370830023853209 +A[10, 4] = 1.09143734899672957818500254654 +A[10, 5] = -8.14978701074692612513997267357 +A[10, 6] = -1.85200656599969598641566180701e1 +A[10, 7] = 2.27394870993505042818970056734e1 +A[10, 8] = 2.49360555267965238987089396762 +A[10, 9] = -3.0467644718982195003823669022 + +A[11, 0] = 2.27331014751653820792359768449 +A[11, 3] = -1.05344954667372501984066689879e1 +A[11, 4] = -2.00087205822486249909675718444 +A[11, 5] = -1.79589318631187989172765950534e1 +A[11, 6] = 2.79488845294199600508499808837e1 +A[11, 7] = -2.85899827713502369474065508674 +A[11, 8] = -8.87285693353062954433549289258 +A[11, 9] = 1.23605671757943030647266201528e1 +A[11, 10] = 6.43392746015763530355970484046e-1 + +A[12, 0] = 5.42937341165687622380535766363e-2 +A[12, 5] = 4.45031289275240888144113950566 +A[12, 6] = 1.89151789931450038304281599044 +A[12, 7] = -5.8012039600105847814672114227 +A[12, 8] = 3.1116436695781989440891606237e-1 +A[12, 9] = -1.52160949662516078556178806805e-1 +A[12, 10] = 2.01365400804030348374776537501e-1 +A[12, 11] = 4.47106157277725905176885569043e-2 + +A[13, 0] = 5.61675022830479523392909219681e-2 +A[13, 6] = 2.53500210216624811088794765333e-1 +A[13, 7] = -2.46239037470802489917441475441e-1 +A[13, 8] = -1.24191423263816360469010140626e-1 +A[13, 9] = 1.5329179827876569731206322685e-1 +A[13, 10] = 8.20105229563468988491666602057e-3 +A[13, 11] = 7.56789766054569976138603589584e-3 +A[13, 12] = -8.298e-3 + +A[14, 0] = 3.18346481635021405060768473261e-2 +A[14, 5] = 2.83009096723667755288322961402e-2 +A[14, 6] = 5.35419883074385676223797384372e-2 +A[14, 7] = -5.49237485713909884646569340306e-2 +A[14, 10] = -1.08347328697249322858509316994e-4 +A[14, 11] = 3.82571090835658412954920192323e-4 +A[14, 12] = -3.40465008687404560802977114492e-4 +A[14, 13] = 1.41312443674632500278074618366e-1 + +A[15, 0] = -4.28896301583791923408573538692e-1 +A[15, 5] = -4.69762141536116384314449447206 +A[15, 6] = 7.68342119606259904184240953878 +A[15, 7] = 4.06898981839711007970213554331 +A[15, 8] = 3.56727187455281109270669543021e-1 +A[15, 12] = -1.39902416515901462129418009734e-3 +A[15, 13] = 2.9475147891527723389556272149 +A[15, 14] = -9.15095847217987001081870187138 + + +B = A[N_STAGES, :N_STAGES] + +E3 = np.zeros(N_STAGES + 1) +E3[:-1] = B.copy() +E3[0] -= 0.244094488188976377952755905512 +E3[8] -= 0.733846688281611857341361741547 +E3[11] -= 0.220588235294117647058823529412e-1 + +E5 = np.zeros(N_STAGES + 1) +E5[0] = 0.1312004499419488073250102996e-1 +E5[5] = -0.1225156446376204440720569753e+1 +E5[6] = -0.4957589496572501915214079952 +E5[7] = 0.1664377182454986536961530415e+1 +E5[8] = -0.3503288487499736816886487290 +E5[9] = 0.3341791187130174790297318841 +E5[10] = 0.8192320648511571246570742613e-1 +E5[11] = -0.2235530786388629525884427845e-1 + +# First 3 coefficients are computed separately. +D = np.zeros((INTERPOLATOR_POWER - 3, N_STAGES_EXTENDED)) +D[0, 0] = -0.84289382761090128651353491142e+1 +D[0, 5] = 0.56671495351937776962531783590 +D[0, 6] = -0.30689499459498916912797304727e+1 +D[0, 7] = 0.23846676565120698287728149680e+1 +D[0, 8] = 0.21170345824450282767155149946e+1 +D[0, 9] = -0.87139158377797299206789907490 +D[0, 10] = 0.22404374302607882758541771650e+1 +D[0, 11] = 0.63157877876946881815570249290 +D[0, 12] = -0.88990336451333310820698117400e-1 +D[0, 13] = 0.18148505520854727256656404962e+2 +D[0, 14] = -0.91946323924783554000451984436e+1 +D[0, 15] = -0.44360363875948939664310572000e+1 + +D[1, 0] = 0.10427508642579134603413151009e+2 +D[1, 5] = 0.24228349177525818288430175319e+3 +D[1, 6] = 0.16520045171727028198505394887e+3 +D[1, 7] = -0.37454675472269020279518312152e+3 +D[1, 8] = -0.22113666853125306036270938578e+2 +D[1, 9] = 0.77334326684722638389603898808e+1 +D[1, 10] = -0.30674084731089398182061213626e+2 +D[1, 11] = -0.93321305264302278729567221706e+1 +D[1, 12] = 0.15697238121770843886131091075e+2 +D[1, 13] = -0.31139403219565177677282850411e+2 +D[1, 14] = -0.93529243588444783865713862664e+1 +D[1, 15] = 0.35816841486394083752465898540e+2 + +D[2, 0] = 0.19985053242002433820987653617e+2 +D[2, 5] = -0.38703730874935176555105901742e+3 +D[2, 6] = -0.18917813819516756882830838328e+3 +D[2, 7] = 0.52780815920542364900561016686e+3 +D[2, 8] = -0.11573902539959630126141871134e+2 +D[2, 9] = 0.68812326946963000169666922661e+1 +D[2, 10] = -0.10006050966910838403183860980e+1 +D[2, 11] = 0.77771377980534432092869265740 +D[2, 12] = -0.27782057523535084065932004339e+1 +D[2, 13] = -0.60196695231264120758267380846e+2 +D[2, 14] = 0.84320405506677161018159903784e+2 +D[2, 15] = 0.11992291136182789328035130030e+2 + +D[3, 0] = -0.25693933462703749003312586129e+2 +D[3, 5] = -0.15418974869023643374053993627e+3 +D[3, 6] = -0.23152937917604549567536039109e+3 +D[3, 7] = 0.35763911791061412378285349910e+3 +D[3, 8] = 0.93405324183624310003907691704e+2 +D[3, 9] = -0.37458323136451633156875139351e+2 +D[3, 10] = 0.10409964950896230045147246184e+3 +D[3, 11] = 0.29840293426660503123344363579e+2 +D[3, 12] = -0.43533456590011143754432175058e+2 +D[3, 13] = 0.96324553959188282948394950600e+2 +D[3, 14] = -0.39177261675615439165231486172e+2 +D[3, 15] = -0.14972683625798562581422125276e+3 diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index e69de29bb..8106d7e0d 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -0,0 +1,496 @@ + +from warnings import warn + +import numpy as np + +from . import _dop853_coefficients as dop853_coefficients +from ._base import DenseOutput, OdeSolver + + +EPS = np.finfo(float).eps + +# Multiply steps computed from asymptotic behaviour of errors by this. +SAFETY = 0.9 + +MIN_FACTOR = 0.2 # Minimum allowed decrease in a step size. +MAX_FACTOR = 10 # Maximum allowed increase in a step size. + + +def norm(x): + """Compute RMS norm.""" + return np.linalg.norm(x) / x.size ** 0.5 + + +def rk_step(fun, t, y, f, h, A, B, C, K): + """Perform a single Runge-Kutta step. + + This function computes a prediction of an explicit Runge-Kutta method and + also estimates the error of a less accurate method. + + Notation for Butcher tableau is as in [1]_. + + Parameters + ---------- + fun : callable + Right-hand side of the system. + t : float + Current time. + y : ndarray, shape (n,) + Current state. + f : ndarray, shape (n,) + Current value of the derivative, i.e., ``fun(x, y)``. + h : float + Step to use. + A : ndarray, shape (n_stages, n_stages) + Coefficients for combining previous RK stages to compute the next + stage. For explicit methods the coefficients at and above the main + diagonal are zeros. + B : ndarray, shape (n_stages,) + Coefficients for combining RK stages for computing the final + prediction. + C : ndarray, shape (n_stages,) + Coefficients for incrementing time for consecutive RK stages. + The value for the first stage is always zero. + K : ndarray, shape (n_stages + 1, n) + Storage array for putting RK stages here. Stages are stored in rows. + The last row is a linear combination of the previous rows with + coefficients + + Returns + ------- + y_new : ndarray, shape (n,) + Solution at t + h computed with a higher accuracy. + f_new : ndarray, shape (n,) + Derivative ``fun(t + h, y_new)``. + + References + ---------- + .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential + Equations I: Nonstiff Problems", Sec. II.4. + """ + K[0] = f + for s, (a, c) in enumerate(zip(A[1:], C[1:]), start=1): + dy = np.dot(K[:s].T, a[:s]) * h + K[s] = fun(t + c * h, y + dy) + + y_new = y + h * np.dot(K[:-1].T, B) + f_new = fun(t + h, y_new) + + K[-1] = f_new + + return y_new, f_new + + +def select_initial_step(fun, t0, y0, f0, direction, order, rtol, atol): + """Empirically select a good initial step. + + The algorithm is described in [1]_. + + Parameters + ---------- + fun : callable + Right-hand side of the system. + t0 : float + Initial value of the independent variable. + y0 : ndarray, shape (n,) + Initial value of the dependent variable. + f0 : ndarray, shape (n,) + Initial value of the derivative, i.e., ``fun(t0, y0)``. + direction : float + Integration direction. + order : float + Error estimator order. It means that the error controlled by the + algorithm is proportional to ``step_size ** (order + 1)`. + rtol : float + Desired relative tolerance. + atol : float + Desired absolute tolerance. + + Returns + ------- + h_abs : float + Absolute value of the suggested initial step. + + References + ---------- + .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential + Equations I: Nonstiff Problems", Sec. II.4. + """ + if y0.size == 0: + return np.inf + + scale = atol + np.abs(y0) * rtol + d0 = norm(y0 / scale) + d1 = norm(f0 / scale) + if d0 < 1e-5 or d1 < 1e-5: + h0 = 1e-6 + else: + h0 = 0.01 * d0 / d1 + + y1 = y0 + h0 * direction * f0 + f1 = fun(t0 + h0 * direction, y1) + d2 = norm((f1 - f0) / scale) / h0 + + if d1 <= 1e-15 and d2 <= 1e-15: + h1 = max(1e-6, h0 * 1e-3) + else: + h1 = (0.01 / max(d1, d2)) ** (1 / (order + 1)) + + return min(100 * h0, h1) + + +def validate_first_step(first_step, t0, t_bound): + """Assert that first_step is valid and return it.""" + if first_step <= 0: + raise ValueError("`first_step` must be positive.") + if first_step > np.abs(t_bound - t0): + raise ValueError("`first_step` exceeds bounds.") + return first_step + + +def validate_max_step(max_step): + """Assert that max_Step is valid and return it.""" + if max_step <= 0: + raise ValueError("`max_step` must be positive.") + return max_step + + +def warn_extraneous(extraneous): + """Display a warning for extraneous keyword arguments. + + The initializer of each solver class is expected to collect keyword + arguments that it doesn't understand and warn about them. This function + prints a warning for each key in the supplied dictionary. + + Parameters + ---------- + extraneous : dict + Extraneous keyword arguments + """ + if extraneous: + warn("The following arguments have no effect for a chosen solver: {}." + .format(", ".join(f"`{x}`" for x in extraneous))) + + +def validate_tol(rtol, atol, n): + """Validate tolerance values.""" + + if np.any(rtol < 100 * EPS): + warn("At least one element of `rtol` is too small. " + f"Setting `rtol = np.maximum(rtol, {100 * EPS})`.") + rtol = np.maximum(rtol, 100 * EPS) + + atol = np.asarray(atol) + if atol.ndim > 0 and atol.shape != (n,): + raise ValueError("`atol` has wrong shape.") + + if np.any(atol < 0): + raise ValueError("`atol` must be positive.") + + return rtol, atol + + +class RungeKutta(OdeSolver): + """Base class for explicit Runge-Kutta methods.""" + C: np.ndarray = NotImplemented + A: np.ndarray = NotImplemented + B: np.ndarray = NotImplemented + E: np.ndarray = NotImplemented + P: np.ndarray = NotImplemented + order: int = NotImplemented + error_estimator_order: int = NotImplemented + n_stages: int = NotImplemented + + def __init__(self, fun, t0, y0, t_bound, max_step=np.inf, + rtol=1e-3, atol=1e-6, vectorized=False, + first_step=None, **extraneous): + warn_extraneous(extraneous) + super().__init__(fun, t0, y0, t_bound, vectorized, + support_complex=True) + self.y_old = None + self.max_step = validate_max_step(max_step) + self.rtol, self.atol = validate_tol(rtol, atol, self.n) + self.f = self.fun(self.t, self.y) + if first_step is None: + self.h_abs = select_initial_step( + self.fun, self.t, self.y, self.f, self.direction, + self.error_estimator_order, self.rtol, self.atol) + else: + self.h_abs = validate_first_step(first_step, t0, t_bound) + self.K = np.empty((self.n_stages + 1, self.n), dtype=self.y.dtype) + self.error_exponent = -1 / (self.error_estimator_order + 1) + self.h_previous = None + + def _estimate_error(self, K, h): + return np.dot(K.T, self.E) * h + + def _estimate_error_norm(self, K, h, scale): + return norm(self._estimate_error(K, h) / scale) + + def _step_impl(self): + t = self.t + y = self.y + + max_step = self.max_step + rtol = self.rtol + atol = self.atol + + min_step = 10 * np.abs(np.nextafter(t, self.direction * np.inf) - t) + + if self.h_abs > max_step: + h_abs = max_step + elif self.h_abs < min_step: + h_abs = min_step + else: + h_abs = self.h_abs + + step_accepted = False + step_rejected = False + + while not step_accepted: + if h_abs < min_step: + return False, self.TOO_SMALL_STEP + + h = h_abs * self.direction + t_new = t + h + + if self.direction * (t_new - self.t_bound) > 0: + t_new = self.t_bound + + h = t_new - t + h_abs = np.abs(h) + + y_new, f_new = rk_step(self.fun, t, y, self.f, h, self.A, + self.B, self.C, self.K) + scale = atol + np.maximum(np.abs(y), np.abs(y_new)) * rtol + error_norm = self._estimate_error_norm(self.K, h, scale) + + if error_norm < 1: + if error_norm == 0: + factor = MAX_FACTOR + else: + factor = min(MAX_FACTOR, + SAFETY * error_norm ** self.error_exponent) + + if step_rejected: + factor = min(1, factor) + + h_abs *= factor + + step_accepted = True + else: + h_abs *= max(MIN_FACTOR, + SAFETY * error_norm ** self.error_exponent) + step_rejected = True + + self.h_previous = h + self.y_old = y + + self.t = t_new + self.y = y_new + + self.h_abs = h_abs + self.f = f_new + + return True, None + + def _dense_output_impl(self): + Q = self.K.T.dot(self.P) + return RkDenseOutput(self.t_old, self.t, self.y_old, Q) + + +class DOP853(RungeKutta): + """Explicit Runge-Kutta method of order 8. + + This is a Python implementation of "DOP853" algorithm originally written + in Fortran [1]_, [2]_. Note that this is not a literate translation, but + the algorithmic core and coefficients are the same. + + Can be applied in the complex domain. + + Parameters + ---------- + fun : callable + Right-hand side of the system. The calling signature is ``fun(t, y)``. + Here, ``t`` is a scalar, and there are two options for the ndarray ``y``: + It can either have shape (n,); then ``fun`` must return array_like with + shape (n,). Alternatively it can have shape (n, k); then ``fun`` + must return an array_like with shape (n, k), i.e. each column + corresponds to a single column in ``y``. The choice between the two + options is determined by `vectorized` argument (see below). + t0 : float + Initial time. + y0 : array_like, shape (n,) + Initial state. + t_bound : float + Boundary time - the integration won't continue beyond it. It also + determines the direction of the integration. + first_step : float or None, optional + Initial step size. Default is ``None`` which means that the algorithm + should choose. + max_step : float, optional + Maximum allowed step size. Default is np.inf, i.e. the step size is not + bounded and determined solely by the solver. + rtol, atol : float and array_like, optional + Relative and absolute tolerances. The solver keeps the local error + estimates less than ``atol + rtol * abs(y)``. Here `rtol` controls a + relative accuracy (number of correct digits), while `atol` controls + absolute accuracy (number of correct decimal places). To achieve the + desired `rtol`, set `atol` to be smaller than the smallest value that + can be expected from ``rtol * abs(y)`` so that `rtol` dominates the + allowable error. If `atol` is larger than ``rtol * abs(y)`` the + number of correct digits is not guaranteed. Conversely, to achieve the + desired `atol` set `rtol` such that ``rtol * abs(y)`` is always smaller + than `atol`. If components of y have different scales, it might be + beneficial to set different `atol` values for different components by + passing array_like with shape (n,) for `atol`. Default values are + 1e-3 for `rtol` and 1e-6 for `atol`. + vectorized : bool, optional + Whether `fun` is implemented in a vectorized fashion. Default is False. + + Attributes + ---------- + n : int + Number of equations. + status : string + Current status of the solver: 'running', 'finished' or 'failed'. + t_bound : float + Boundary time. + direction : float + Integration direction: +1 or -1. + t : float + Current time. + y : ndarray + Current state. + t_old : float + Previous time. None if no steps were made yet. + step_size : float + Size of the last successful step. None if no steps were made yet. + nfev : int + Number evaluations of the system's right-hand side. + njev : int + Number of evaluations of the Jacobian. Is always 0 for this solver + as it does not use the Jacobian. + nlu : int + Number of LU decompositions. Is always 0 for this solver. + + References + ---------- + .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential + Equations I: Nonstiff Problems", Sec. II. + .. [2] `Page with original Fortran code of DOP853 + `_. + """ + n_stages = dop853_coefficients.N_STAGES + order = 8 + error_estimator_order = 7 + A = dop853_coefficients.A[:n_stages, :n_stages] + B = dop853_coefficients.B + C = dop853_coefficients.C[:n_stages] + E3 = dop853_coefficients.E3 + E5 = dop853_coefficients.E5 + D = dop853_coefficients.D + + A_EXTRA = dop853_coefficients.A[n_stages + 1:] + C_EXTRA = dop853_coefficients.C[n_stages + 1:] + + def __init__(self, fun, t0, y0, t_bound, max_step=np.inf, + rtol=1e-3, atol=1e-6, vectorized=False, + first_step=None, **extraneous): + super().__init__(fun, t0, y0, t_bound, max_step, rtol, atol, + vectorized, first_step, **extraneous) + self.K_extended = np.empty((dop853_coefficients.N_STAGES_EXTENDED, + self.n), dtype=self.y.dtype) + self.K = self.K_extended[:self.n_stages + 1] + + def _estimate_error(self, K, h): # Left for testing purposes. + err5 = np.dot(K.T, self.E5) + err3 = np.dot(K.T, self.E3) + denom = np.hypot(np.abs(err5), 0.1 * np.abs(err3)) + correction_factor = np.ones_like(err5) + mask = denom > 0 + correction_factor[mask] = np.abs(err5[mask]) / denom[mask] + return h * err5 * correction_factor + + def _estimate_error_norm(self, K, h, scale): + err5 = np.dot(K.T, self.E5) / scale + err3 = np.dot(K.T, self.E3) / scale + err5_norm_2 = np.linalg.norm(err5)**2 + err3_norm_2 = np.linalg.norm(err3)**2 + if err5_norm_2 == 0 and err3_norm_2 == 0: + return 0.0 + denom = err5_norm_2 + 0.01 * err3_norm_2 + return np.abs(h) * err5_norm_2 / np.sqrt(denom * len(scale)) + + def _dense_output_impl(self): + K = self.K_extended + h = self.h_previous + for s, (a, c) in enumerate(zip(self.A_EXTRA, self.C_EXTRA), + start=self.n_stages + 1): + dy = np.dot(K[:s].T, a[:s]) * h + K[s] = self.fun(self.t_old + c * h, self.y_old + dy) + + F = np.empty((dop853_coefficients.INTERPOLATOR_POWER, self.n), + dtype=self.y_old.dtype) + + f_old = K[0] + delta_y = self.y - self.y_old + + F[0] = delta_y + F[1] = h * f_old - delta_y + F[2] = 2 * delta_y - h * (self.f + f_old) + F[3:] = h * np.dot(self.D, K) + + return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) + + +class Dop853DenseOutput(DenseOutput): + def __init__(self, t_old, t, y_old, F): + super().__init__(t_old, t) + self.h = t - t_old + self.F = F + self.y_old = y_old + + def _call_impl(self, t): + x = (t - self.t_old) / self.h + + if t.ndim == 0: + y = np.zeros_like(self.y_old) + else: + x = x[:, None] + y = np.zeros((len(x), len(self.y_old)), dtype=self.y_old.dtype) + + for i, f in enumerate(reversed(self.F)): + y += f + if i % 2 == 0: + y *= x + else: + y *= 1 - x + y += self.y_old + + return y.T + + +class RkDenseOutput(DenseOutput): + def __init__(self, t_old, t, y_old, Q): + super().__init__(t_old, t) + self.h = t - t_old + self.Q = Q + self.order = Q.shape[1] - 1 + self.y_old = y_old + + def _call_impl(self, t): + x = (t - self.t_old) / self.h + if t.ndim == 0: + p = np.tile(x, self.order + 1) + p = np.cumprod(p) + else: + p = np.tile(x, (self.order + 1, 1)) + p = np.cumprod(p, axis=0) + y = self.h * np.dot(self.Q, p) + if y.ndim == 2: + y += self.y_old[:, None] + else: + y += self.y_old + + return y From 942df193f0251344ed8fb78af791755d0b478b98 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 9 Jan 2024 13:57:17 +0100 Subject: [PATCH 080/346] more ivp draft --- src/hapsira/core/math/ivp/_api.py | 7 +- src/hapsira/core/math/ivp/_base.py | 33 +++-- src/hapsira/core/math/ivp/_brentq.py | 48 ++++--- src/hapsira/core/math/ivp/_common.py | 16 +-- .../core/math/ivp/_dop853_coefficients.py | 132 +++++++++--------- src/hapsira/core/math/ivp/_rk.py | 116 ++++++++++----- 6 files changed, 206 insertions(+), 146 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index e7a74627a..b01bb8e30 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -1,4 +1,3 @@ - import numpy as np from ._brentq import brentq @@ -7,8 +6,10 @@ from ._rk import EPS, DOP853 -MESSAGES = {0: "The solver successfully reached the end of the integration interval.", - 1: "A termination event occurred."} +MESSAGES = { + 0: "The solver successfully reached the end of the integration interval.", + 1: "A termination event occurred.", +} class OdeResult(dict): diff --git a/src/hapsira/core/math/ivp/_base.py b/src/hapsira/core/math/ivp/_base.py index 05c00e60c..a19084e33 100644 --- a/src/hapsira/core/math/ivp/_base.py +++ b/src/hapsira/core/math/ivp/_base.py @@ -1,4 +1,3 @@ - import numpy as np @@ -7,8 +6,10 @@ def check_arguments(fun, y0, support_complex): y0 = np.asarray(y0) if np.issubdtype(y0.dtype, np.complexfloating): if not support_complex: - raise ValueError("`y0` is complex, but the chosen solver does " - "not support integration in a complex domain.") + raise ValueError( + "`y0` is complex, but the chosen solver does " + "not support integration in a complex domain." + ) dtype = complex else: dtype = float @@ -127,10 +128,10 @@ class OdeSolver: nlu : int Number of LU decompositions. """ + TOO_SMALL_STEP = "Required step size is less than spacing between numbers." - def __init__(self, fun, t0, y0, t_bound, vectorized, - support_complex=False): + def __init__(self, fun, t0, y0, t_bound, vectorized, support_complex=False): self.t_old = None self.t = t0 self._fun, self.y = check_arguments(fun, y0, support_complex) @@ -138,8 +139,10 @@ def __init__(self, fun, t0, y0, t_bound, vectorized, self.vectorized = vectorized if vectorized: + def fun_single(t, y): return self._fun(t, y[:, None]).ravel() + fun_vectorized = self._fun else: fun_single = self._fun @@ -160,7 +163,7 @@ def fun(t, y): self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 self.n = self.y.size - self.status = 'running' + self.status = "running" self.nfev = 0 self.njev = 0 @@ -183,26 +186,25 @@ def step(self): `self.status` is 'failed' after the step was taken or None otherwise. """ - if self.status != 'running': - raise RuntimeError("Attempt to step on a failed or finished " - "solver.") + if self.status != "running": + raise RuntimeError("Attempt to step on a failed or finished " "solver.") if self.n == 0 or self.t == self.t_bound: # Handle corner cases of empty solver or no integration. self.t_old = self.t self.t = self.t_bound message = None - self.status = 'finished' + self.status = "finished" else: t = self.t success, message = self._step_impl() if not success: - self.status = 'failed' + self.status = "failed" else: self.t_old = t if self.direction * (self.t - self.t_bound) >= 0: - self.status = 'finished' + self.status = "finished" return message @@ -215,8 +217,9 @@ def dense_output(self): Local interpolant over the last successful step. """ if self.t_old is None: - raise RuntimeError("Dense output is available after a successful " - "step was made.") + raise RuntimeError( + "Dense output is available after a successful " "step was made." + ) if self.n == 0 or self.t == self.t_old: # Handle corner cases of empty solver and no integration. @@ -243,6 +246,7 @@ class DenseOutput: t_min, t_max : float Time range of the interpolation. """ + def __init__(self, t_old, t): self.t_old = t_old self.t = t @@ -278,6 +282,7 @@ class ConstantDenseOutput(DenseOutput): This class used for degenerate integration cases: equal integration limits or a system with 0 equations. """ + def __init__(self, t_old, t, value): super().__init__(t_old, t) self.value = value diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index e7e072b6e..50993f4db 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -1,4 +1,3 @@ - import operator import numpy as np @@ -13,18 +12,23 @@ _EVALUEERR = -3 _EINPROGRESS = 1 -CONVERGED = 'converged' -SIGNERR = 'sign error' -CONVERR = 'convergence error' -VALUEERR = 'value error' -INPROGRESS = 'No error' +CONVERGED = "converged" +SIGNERR = "sign error" +CONVERR = "convergence error" +VALUEERR = "value error" +INPROGRESS = "No error" -flag_map = {_ECONVERGED: CONVERGED, _ESIGNERR: SIGNERR, _ECONVERR: CONVERR, - _EVALUEERR: VALUEERR, _EINPROGRESS: INPROGRESS} +flag_map = { + _ECONVERGED: CONVERGED, + _ESIGNERR: SIGNERR, + _ECONVERR: CONVERR, + _EVALUEERR: VALUEERR, + _EINPROGRESS: INPROGRESS, +} class OptimizeResult(dict): - """ Represents the optimization result. + """Represents the optimization result. Attributes ---------- @@ -134,13 +138,11 @@ def __init__(self, root, iterations, function_calls, flag): def _wrap_nan_raise(f): - def f_raise(x, *args): fx = f(x, *args) f_raise._function_calls += 1 if np.isnan(fx): - msg = (f'The function value at x={x} is NaN; ' - 'solver cannot continue.') + msg = f"The function value at x={x} is NaN; " "solver cannot continue." err = ValueError(msg) err._x = x err._function_calls = f_raise._function_calls @@ -154,10 +156,9 @@ def f_raise(x, *args): def results_c(full_output, r): if full_output: x, funcalls, iterations, flag = r - results = RootResults(root=x, - iterations=iterations, - function_calls=funcalls, - flag=flag) + results = RootResults( + root=x, iterations=iterations, function_calls=funcalls, flag=flag + ) return x, results else: return r @@ -168,9 +169,17 @@ def results_c(full_output, r): _rtol = 4 * np.finfo(float).eps -def brentq(f, a, b, args=(), - xtol=_xtol, rtol=_rtol, maxiter=_iter, - full_output=False, disp=True): +def brentq( + f, + a, + b, + args=(), + xtol=_xtol, + rtol=_rtol, + maxiter=_iter, + full_output=False, + disp=True, +): """ Find a root of a function in a bracketing interval using Brent's method. @@ -297,4 +306,3 @@ def brentq(f, a, b, args=(), f = _wrap_nan_raise(f) r = _zeros._brentq(f, a, b, xtol, rtol, maxiter, args, full_output, disp) return results_c(full_output, r) - diff --git a/src/hapsira/core/math/ivp/_common.py b/src/hapsira/core/math/ivp/_common.py index ad3bf5a45..692e7eb71 100644 --- a/src/hapsira/core/math/ivp/_common.py +++ b/src/hapsira/core/math/ivp/_common.py @@ -1,4 +1,3 @@ - from itertools import groupby import numpy as np @@ -33,18 +32,17 @@ class OdeSolution: t_min, t_max : float Time range of the interpolation. """ + def __init__(self, ts, interpolants): ts = np.asarray(ts) d = np.diff(ts) # The first case covers integration on zero segment. - if not ((ts.size == 2 and ts[0] == ts[-1]) - or np.all(d > 0) or np.all(d < 0)): + if not ((ts.size == 2 and ts[0] == ts[-1]) or np.all(d > 0) or np.all(d < 0)): raise ValueError("`ts` must be strictly increasing or decreasing.") self.n_segments = len(interpolants) if ts.shape != (self.n_segments + 1,): - raise ValueError("Numbers of time stamps and interpolants " - "don't match.") + raise ValueError("Numbers of time stamps and interpolants " "don't match.") self.ts = ts self.interpolants = interpolants @@ -63,9 +61,9 @@ def _call_single(self, t): # Here we preserve a certain symmetry that when t is in self.ts, # then we prioritize a segment with a lower index. if self.ascending: - ind = np.searchsorted(self.ts_sorted, t, side='left') + ind = np.searchsorted(self.ts_sorted, t, side="left") else: - ind = np.searchsorted(self.ts_sorted, t, side='right') + ind = np.searchsorted(self.ts_sorted, t, side="right") segment = min(max(ind - 1, 0), self.n_segments - 1) if not self.ascending: @@ -99,9 +97,9 @@ def __call__(self, t): # See comment in self._call_single. if self.ascending: - segments = np.searchsorted(self.ts_sorted, t_sorted, side='left') + segments = np.searchsorted(self.ts_sorted, t_sorted, side="left") else: - segments = np.searchsorted(self.ts_sorted, t_sorted, side='right') + segments = np.searchsorted(self.ts_sorted, t_sorted, side="right") segments -= 1 segments[segments < 0] = 0 segments[segments > self.n_segments - 1] = self.n_segments - 1 diff --git a/src/hapsira/core/math/ivp/_dop853_coefficients.py b/src/hapsira/core/math/ivp/_dop853_coefficients.py index f39f2f365..7f9d03aff 100644 --- a/src/hapsira/core/math/ivp/_dop853_coefficients.py +++ b/src/hapsira/core/math/ivp/_dop853_coefficients.py @@ -4,22 +4,26 @@ N_STAGES_EXTENDED = 16 INTERPOLATOR_POWER = 7 -C = np.array([0.0, - 0.526001519587677318785587544488e-01, - 0.789002279381515978178381316732e-01, - 0.118350341907227396726757197510, - 0.281649658092772603273242802490, - 0.333333333333333333333333333333, - 0.25, - 0.307692307692307692307692307692, - 0.651282051282051282051282051282, - 0.6, - 0.857142857142857142857142857142, - 1.0, - 1.0, - 0.1, - 0.2, - 0.777777777777777777777777777778]) +C = np.array( + [ + 0.0, + 0.526001519587677318785587544488e-01, + 0.789002279381515978178381316732e-01, + 0.118350341907227396726757197510, + 0.281649658092772603273242802490, + 0.333333333333333333333333333333, + 0.25, + 0.307692307692307692307692307692, + 0.651282051282051282051282051282, + 0.6, + 0.857142857142857142857142857142, + 1.0, + 1.0, + 0.1, + 0.2, + 0.777777777777777777777777777778, + ] +) A = np.zeros((N_STAGES_EXTENDED, N_STAGES_EXTENDED)) A[1, 0] = 5.26001519587677318785587544488e-2 @@ -130,9 +134,9 @@ E5 = np.zeros(N_STAGES + 1) E5[0] = 0.1312004499419488073250102996e-1 -E5[5] = -0.1225156446376204440720569753e+1 +E5[5] = -0.1225156446376204440720569753e1 E5[6] = -0.4957589496572501915214079952 -E5[7] = 0.1664377182454986536961530415e+1 +E5[7] = 0.1664377182454986536961530415e1 E5[8] = -0.3503288487499736816886487290 E5[9] = 0.3341791187130174790297318841 E5[10] = 0.8192320648511571246570742613e-1 @@ -140,54 +144,54 @@ # First 3 coefficients are computed separately. D = np.zeros((INTERPOLATOR_POWER - 3, N_STAGES_EXTENDED)) -D[0, 0] = -0.84289382761090128651353491142e+1 +D[0, 0] = -0.84289382761090128651353491142e1 D[0, 5] = 0.56671495351937776962531783590 -D[0, 6] = -0.30689499459498916912797304727e+1 -D[0, 7] = 0.23846676565120698287728149680e+1 -D[0, 8] = 0.21170345824450282767155149946e+1 +D[0, 6] = -0.30689499459498916912797304727e1 +D[0, 7] = 0.23846676565120698287728149680e1 +D[0, 8] = 0.21170345824450282767155149946e1 D[0, 9] = -0.87139158377797299206789907490 -D[0, 10] = 0.22404374302607882758541771650e+1 +D[0, 10] = 0.22404374302607882758541771650e1 D[0, 11] = 0.63157877876946881815570249290 D[0, 12] = -0.88990336451333310820698117400e-1 -D[0, 13] = 0.18148505520854727256656404962e+2 -D[0, 14] = -0.91946323924783554000451984436e+1 -D[0, 15] = -0.44360363875948939664310572000e+1 - -D[1, 0] = 0.10427508642579134603413151009e+2 -D[1, 5] = 0.24228349177525818288430175319e+3 -D[1, 6] = 0.16520045171727028198505394887e+3 -D[1, 7] = -0.37454675472269020279518312152e+3 -D[1, 8] = -0.22113666853125306036270938578e+2 -D[1, 9] = 0.77334326684722638389603898808e+1 -D[1, 10] = -0.30674084731089398182061213626e+2 -D[1, 11] = -0.93321305264302278729567221706e+1 -D[1, 12] = 0.15697238121770843886131091075e+2 -D[1, 13] = -0.31139403219565177677282850411e+2 -D[1, 14] = -0.93529243588444783865713862664e+1 -D[1, 15] = 0.35816841486394083752465898540e+2 - -D[2, 0] = 0.19985053242002433820987653617e+2 -D[2, 5] = -0.38703730874935176555105901742e+3 -D[2, 6] = -0.18917813819516756882830838328e+3 -D[2, 7] = 0.52780815920542364900561016686e+3 -D[2, 8] = -0.11573902539959630126141871134e+2 -D[2, 9] = 0.68812326946963000169666922661e+1 -D[2, 10] = -0.10006050966910838403183860980e+1 +D[0, 13] = 0.18148505520854727256656404962e2 +D[0, 14] = -0.91946323924783554000451984436e1 +D[0, 15] = -0.44360363875948939664310572000e1 + +D[1, 0] = 0.10427508642579134603413151009e2 +D[1, 5] = 0.24228349177525818288430175319e3 +D[1, 6] = 0.16520045171727028198505394887e3 +D[1, 7] = -0.37454675472269020279518312152e3 +D[1, 8] = -0.22113666853125306036270938578e2 +D[1, 9] = 0.77334326684722638389603898808e1 +D[1, 10] = -0.30674084731089398182061213626e2 +D[1, 11] = -0.93321305264302278729567221706e1 +D[1, 12] = 0.15697238121770843886131091075e2 +D[1, 13] = -0.31139403219565177677282850411e2 +D[1, 14] = -0.93529243588444783865713862664e1 +D[1, 15] = 0.35816841486394083752465898540e2 + +D[2, 0] = 0.19985053242002433820987653617e2 +D[2, 5] = -0.38703730874935176555105901742e3 +D[2, 6] = -0.18917813819516756882830838328e3 +D[2, 7] = 0.52780815920542364900561016686e3 +D[2, 8] = -0.11573902539959630126141871134e2 +D[2, 9] = 0.68812326946963000169666922661e1 +D[2, 10] = -0.10006050966910838403183860980e1 D[2, 11] = 0.77771377980534432092869265740 -D[2, 12] = -0.27782057523535084065932004339e+1 -D[2, 13] = -0.60196695231264120758267380846e+2 -D[2, 14] = 0.84320405506677161018159903784e+2 -D[2, 15] = 0.11992291136182789328035130030e+2 - -D[3, 0] = -0.25693933462703749003312586129e+2 -D[3, 5] = -0.15418974869023643374053993627e+3 -D[3, 6] = -0.23152937917604549567536039109e+3 -D[3, 7] = 0.35763911791061412378285349910e+3 -D[3, 8] = 0.93405324183624310003907691704e+2 -D[3, 9] = -0.37458323136451633156875139351e+2 -D[3, 10] = 0.10409964950896230045147246184e+3 -D[3, 11] = 0.29840293426660503123344363579e+2 -D[3, 12] = -0.43533456590011143754432175058e+2 -D[3, 13] = 0.96324553959188282948394950600e+2 -D[3, 14] = -0.39177261675615439165231486172e+2 -D[3, 15] = -0.14972683625798562581422125276e+3 +D[2, 12] = -0.27782057523535084065932004339e1 +D[2, 13] = -0.60196695231264120758267380846e2 +D[2, 14] = 0.84320405506677161018159903784e2 +D[2, 15] = 0.11992291136182789328035130030e2 + +D[3, 0] = -0.25693933462703749003312586129e2 +D[3, 5] = -0.15418974869023643374053993627e3 +D[3, 6] = -0.23152937917604549567536039109e3 +D[3, 7] = 0.35763911791061412378285349910e3 +D[3, 8] = 0.93405324183624310003907691704e2 +D[3, 9] = -0.37458323136451633156875139351e2 +D[3, 10] = 0.10409964950896230045147246184e3 +D[3, 11] = 0.29840293426660503123344363579e2 +D[3, 12] = -0.43533456590011143754432175058e2 +D[3, 13] = 0.96324553959188282948394950600e2 +D[3, 14] = -0.39177261675615439165231486172e2 +D[3, 15] = -0.14972683625798562581422125276e3 diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 8106d7e0d..206558596 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -1,4 +1,3 @@ - from warnings import warn import numpy as np @@ -18,7 +17,7 @@ def norm(x): """Compute RMS norm.""" - return np.linalg.norm(x) / x.size ** 0.5 + return np.linalg.norm(x) / x.size**0.5 def rk_step(fun, t, y, f, h, A, B, C, K): @@ -168,16 +167,21 @@ def warn_extraneous(extraneous): Extraneous keyword arguments """ if extraneous: - warn("The following arguments have no effect for a chosen solver: {}." - .format(", ".join(f"`{x}`" for x in extraneous))) + warn( + "The following arguments have no effect for a chosen solver: {}.".format( + ", ".join(f"`{x}`" for x in extraneous) + ) + ) def validate_tol(rtol, atol, n): """Validate tolerance values.""" if np.any(rtol < 100 * EPS): - warn("At least one element of `rtol` is too small. " - f"Setting `rtol = np.maximum(rtol, {100 * EPS})`.") + warn( + "At least one element of `rtol` is too small. " + f"Setting `rtol = np.maximum(rtol, {100 * EPS})`." + ) rtol = np.maximum(rtol, 100 * EPS) atol = np.asarray(atol) @@ -192,6 +196,7 @@ def validate_tol(rtol, atol, n): class RungeKutta(OdeSolver): """Base class for explicit Runge-Kutta methods.""" + C: np.ndarray = NotImplemented A: np.ndarray = NotImplemented B: np.ndarray = NotImplemented @@ -201,20 +206,36 @@ class RungeKutta(OdeSolver): error_estimator_order: int = NotImplemented n_stages: int = NotImplemented - def __init__(self, fun, t0, y0, t_bound, max_step=np.inf, - rtol=1e-3, atol=1e-6, vectorized=False, - first_step=None, **extraneous): + def __init__( + self, + fun, + t0, + y0, + t_bound, + max_step=np.inf, + rtol=1e-3, + atol=1e-6, + vectorized=False, + first_step=None, + **extraneous, + ): warn_extraneous(extraneous) - super().__init__(fun, t0, y0, t_bound, vectorized, - support_complex=True) + super().__init__(fun, t0, y0, t_bound, vectorized, support_complex=True) self.y_old = None self.max_step = validate_max_step(max_step) self.rtol, self.atol = validate_tol(rtol, atol, self.n) self.f = self.fun(self.t, self.y) if first_step is None: self.h_abs = select_initial_step( - self.fun, self.t, self.y, self.f, self.direction, - self.error_estimator_order, self.rtol, self.atol) + self.fun, + self.t, + self.y, + self.f, + self.direction, + self.error_estimator_order, + self.rtol, + self.atol, + ) else: self.h_abs = validate_first_step(first_step, t0, t_bound) self.K = np.empty((self.n_stages + 1, self.n), dtype=self.y.dtype) @@ -260,8 +281,9 @@ def _step_impl(self): h = t_new - t h_abs = np.abs(h) - y_new, f_new = rk_step(self.fun, t, y, self.f, h, self.A, - self.B, self.C, self.K) + y_new, f_new = rk_step( + self.fun, t, y, self.f, h, self.A, self.B, self.C, self.K + ) scale = atol + np.maximum(np.abs(y), np.abs(y_new)) * rtol error_norm = self._estimate_error_norm(self.K, h, scale) @@ -269,8 +291,7 @@ def _step_impl(self): if error_norm == 0: factor = MAX_FACTOR else: - factor = min(MAX_FACTOR, - SAFETY * error_norm ** self.error_exponent) + factor = min(MAX_FACTOR, SAFETY * error_norm**self.error_exponent) if step_rejected: factor = min(1, factor) @@ -279,8 +300,7 @@ def _step_impl(self): step_accepted = True else: - h_abs *= max(MIN_FACTOR, - SAFETY * error_norm ** self.error_exponent) + h_abs *= max(MIN_FACTOR, SAFETY * error_norm**self.error_exponent) step_rejected = True self.h_previous = h @@ -381,6 +401,7 @@ class DOP853(RungeKutta): .. [2] `Page with original Fortran code of DOP853 `_. """ + n_stages = dop853_coefficients.N_STAGES order = 8 error_estimator_order = 7 @@ -391,17 +412,38 @@ class DOP853(RungeKutta): E5 = dop853_coefficients.E5 D = dop853_coefficients.D - A_EXTRA = dop853_coefficients.A[n_stages + 1:] - C_EXTRA = dop853_coefficients.C[n_stages + 1:] - - def __init__(self, fun, t0, y0, t_bound, max_step=np.inf, - rtol=1e-3, atol=1e-6, vectorized=False, - first_step=None, **extraneous): - super().__init__(fun, t0, y0, t_bound, max_step, rtol, atol, - vectorized, first_step, **extraneous) - self.K_extended = np.empty((dop853_coefficients.N_STAGES_EXTENDED, - self.n), dtype=self.y.dtype) - self.K = self.K_extended[:self.n_stages + 1] + A_EXTRA = dop853_coefficients.A[n_stages + 1 :] + C_EXTRA = dop853_coefficients.C[n_stages + 1 :] + + def __init__( + self, + fun, + t0, + y0, + t_bound, + max_step=np.inf, + rtol=1e-3, + atol=1e-6, + vectorized=False, + first_step=None, + **extraneous, + ): + super().__init__( + fun, + t0, + y0, + t_bound, + max_step, + rtol, + atol, + vectorized, + first_step, + **extraneous, + ) + self.K_extended = np.empty( + (dop853_coefficients.N_STAGES_EXTENDED, self.n), dtype=self.y.dtype + ) + self.K = self.K_extended[: self.n_stages + 1] def _estimate_error(self, K, h): # Left for testing purposes. err5 = np.dot(K.T, self.E5) @@ -415,8 +457,8 @@ def _estimate_error(self, K, h): # Left for testing purposes. def _estimate_error_norm(self, K, h, scale): err5 = np.dot(K.T, self.E5) / scale err3 = np.dot(K.T, self.E3) / scale - err5_norm_2 = np.linalg.norm(err5)**2 - err3_norm_2 = np.linalg.norm(err3)**2 + err5_norm_2 = np.linalg.norm(err5) ** 2 + err3_norm_2 = np.linalg.norm(err3) ** 2 if err5_norm_2 == 0 and err3_norm_2 == 0: return 0.0 denom = err5_norm_2 + 0.01 * err3_norm_2 @@ -425,13 +467,15 @@ def _estimate_error_norm(self, K, h, scale): def _dense_output_impl(self): K = self.K_extended h = self.h_previous - for s, (a, c) in enumerate(zip(self.A_EXTRA, self.C_EXTRA), - start=self.n_stages + 1): + for s, (a, c) in enumerate( + zip(self.A_EXTRA, self.C_EXTRA), start=self.n_stages + 1 + ): dy = np.dot(K[:s].T, a[:s]) * h K[s] = self.fun(self.t_old + c * h, self.y_old + dy) - F = np.empty((dop853_coefficients.INTERPOLATOR_POWER, self.n), - dtype=self.y_old.dtype) + F = np.empty( + (dop853_coefficients.INTERPOLATOR_POWER, self.n), dtype=self.y_old.dtype + ) f_old = K[0] delta_y = self.y - self.y_old From d63448d0f0b37669f1237fe7d596a1dd4e85083b Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 9 Jan 2024 14:16:09 +0100 Subject: [PATCH 081/346] more ivp draft --- src/hapsira/core/math/ivp/_brentq.c | 130 +++++++++++++++++++++++++++ src/hapsira/core/math/ivp/_zeros.c | 132 ++++++++++++++++++++++++++++ src/hapsira/core/math/ivp/_zeros.h | 40 +++++++++ 3 files changed, 302 insertions(+) create mode 100644 src/hapsira/core/math/ivp/_brentq.c create mode 100644 src/hapsira/core/math/ivp/_zeros.c create mode 100644 src/hapsira/core/math/ivp/_zeros.h diff --git a/src/hapsira/core/math/ivp/_brentq.c b/src/hapsira/core/math/ivp/_brentq.c new file mode 100644 index 000000000..1d8df1eb8 --- /dev/null +++ b/src/hapsira/core/math/ivp/_brentq.c @@ -0,0 +1,130 @@ +/* Written by Charles Harris charles.harris@sdl.usu.edu */ + +#include +#include "zeros.h" + +#define MIN(a, b) ((a) < (b) ? (a) : (b)) + +/* + At the top of the loop the situation is the following: + + 1. the root is bracketed between xa and xb + 2. xa is the most recent estimate + 3. xp is the previous estimate + 4. |fp| < |fb| + + The order of xa and xp doesn't matter, but assume xp < xb. Then xa lies to + the right of xp and the assumption is that xa is increasing towards the root. + In this situation we will attempt quadratic extrapolation as long as the + condition + + * |fa| < |fp| < |fb| + + is satisfied. That is, the function value is decreasing as we go along. + Note the 4 above implies that the right inequlity already holds. + + The first check is that xa is still to the left of the root. If not, xb is + replaced by xp and the interval reverses, with xb < xa. In this situation + we will try linear interpolation. That this has happened is signaled by the + equality xb == xp; + + The second check is that |fa| < |fb|. If this is not the case, we swap + xa and xb and resort to bisection. + +*/ + +double +brentq(callback_type f, double xa, double xb, double xtol, double rtol, + int iter, void *func_data_param, scipy_zeros_info *solver_stats) +{ + double xpre = xa, xcur = xb; + double xblk = 0., fpre, fcur, fblk = 0., spre = 0., scur = 0., sbis; + /* the tolerance is 2*delta */ + double delta; + double stry, dpre, dblk; + int i; + solver_stats->error_num = INPROGRESS; + + fpre = (*f)(xpre, func_data_param); + fcur = (*f)(xcur, func_data_param); + solver_stats->funcalls = 2; + if (fpre == 0) { + solver_stats->error_num = CONVERGED; + return xpre; + } + if (fcur == 0) { + solver_stats->error_num = CONVERGED; + return xcur; + } + if (signbit(fpre)==signbit(fcur)) { + solver_stats->error_num = SIGNERR; + return 0.; + } + solver_stats->iterations = 0; + for (i = 0; i < iter; i++) { + solver_stats->iterations++; + if (fpre != 0 && fcur != 0 && + (signbit(fpre) != signbit(fcur))) { + xblk = xpre; + fblk = fpre; + spre = scur = xcur - xpre; + } + if (fabs(fblk) < fabs(fcur)) { + xpre = xcur; + xcur = xblk; + xblk = xpre; + + fpre = fcur; + fcur = fblk; + fblk = fpre; + } + + delta = (xtol + rtol*fabs(xcur))/2; + sbis = (xblk - xcur)/2; + if (fcur == 0 || fabs(sbis) < delta) { + solver_stats->error_num = CONVERGED; + return xcur; + } + + if (fabs(spre) > delta && fabs(fcur) < fabs(fpre)) { + if (xpre == xblk) { + /* interpolate */ + stry = -fcur*(xcur - xpre)/(fcur - fpre); + } + else { + /* extrapolate */ + dpre = (fpre - fcur)/(xpre - xcur); + dblk = (fblk - fcur)/(xblk - xcur); + stry = -fcur*(fblk*dblk - fpre*dpre) + /(dblk*dpre*(fblk - fpre)); + } + if (2*fabs(stry) < MIN(fabs(spre), 3*fabs(sbis) - delta)) { + /* good short step */ + spre = scur; + scur = stry; + } else { + /* bisect */ + spre = sbis; + scur = sbis; + } + } + else { + /* bisect */ + spre = sbis; + scur = sbis; + } + + xpre = xcur; fpre = fcur; + if (fabs(scur) > delta) { + xcur += scur; + } + else { + xcur += (sbis > 0 ? delta : -delta); + } + + fcur = (*f)(xcur, func_data_param); + solver_stats->funcalls++; + } + solver_stats->error_num = CONVERR; + return xcur; +} diff --git a/src/hapsira/core/math/ivp/_zeros.c b/src/hapsira/core/math/ivp/_zeros.c new file mode 100644 index 000000000..0516ffab9 --- /dev/null +++ b/src/hapsira/core/math/ivp/_zeros.c @@ -0,0 +1,132 @@ +/* + * Helper function that calls a Python function with extended arguments + */ + +static PyObject * +call_solver(solver_type solver, PyObject *self, PyObject *args) +{ + double a, b, xtol, rtol, zero; + int iter, fulloutput, disp=1, flag=0; + scipy_zeros_parameters params; + scipy_zeros_info solver_stats; + PyObject *f, *xargs; + + if (!PyArg_ParseTuple(args, "OddddiOi|i", + &f, &a, &b, &xtol, &rtol, &iter, &xargs, &fulloutput, &disp)) { + PyErr_SetString(PyExc_RuntimeError, "Unable to parse arguments"); + return NULL; + } + if (xtol < 0) { + PyErr_SetString(PyExc_ValueError, "xtol must be >= 0"); + return NULL; + } + if (iter < 0) { + PyErr_SetString(PyExc_ValueError, "maxiter should be > 0"); + return NULL; + } + + params.function = f; + params.xargs = xargs; + + if (!setjmp(params.env)) { + /* direct return */ + solver_stats.error_num = 0; + zero = solver(scipy_zeros_functions_func, a, b, xtol, rtol, + iter, (void*)¶ms, &solver_stats); + } else { + /* error return from Python function */ + return NULL; + } + + if (solver_stats.error_num != CONVERGED) { + if (solver_stats.error_num == SIGNERR) { + PyErr_SetString(PyExc_ValueError, + "f(a) and f(b) must have different signs"); + return NULL; + } + if (solver_stats.error_num == CONVERR) { + if (disp) { + char msg[100]; + PyOS_snprintf(msg, sizeof(msg), + "Failed to converge after %d iterations.", + solver_stats.iterations); + PyErr_SetString(PyExc_RuntimeError, msg); + return NULL; + } + flag = CONVERR; + } + } + else { + flag = CONVERGED; + } + if (fulloutput) { + return Py_BuildValue("diii", + zero, solver_stats.funcalls, solver_stats.iterations, flag); + } + else { + return Py_BuildValue("d", zero); + } +} + +/* + * These routines interface with the solvers through call_solver + */ + +static PyObject * +_bisect(PyObject *self, PyObject *args) +{ + return call_solver(bisect,self,args); +} + +static PyObject * +_ridder(PyObject *self, PyObject *args) +{ + return call_solver(ridder,self,args); +} + +static PyObject * +_brenth(PyObject *self, PyObject *args) +{ + return call_solver(brenth,self,args); +} + +static PyObject * +_brentq(PyObject *self, PyObject *args) +{ + return call_solver(brentq,self,args); +} + +/* + * Standard Python module interface + */ + +static PyMethodDef +Zerosmethods[] = { + {"_bisect", _bisect, METH_VARARGS, "a"}, + {"_ridder", _ridder, METH_VARARGS, "a"}, + {"_brenth", _brenth, METH_VARARGS, "a"}, + {"_brentq", _brentq, METH_VARARGS, "a"}, + {NULL, NULL} +}; + +static struct PyModuleDef moduledef = { + PyModuleDef_HEAD_INIT, + "_zeros", + NULL, + -1, + Zerosmethods, + NULL, + NULL, + NULL, + NULL +}; + +PyMODINIT_FUNC +PyInit__zeros(void) +{ + PyObject *m; + + m = PyModule_Create(&moduledef); + + return m; +} diff --git a/src/hapsira/core/math/ivp/_zeros.h b/src/hapsira/core/math/ivp/_zeros.h new file mode 100644 index 000000000..14afcc13f --- /dev/null +++ b/src/hapsira/core/math/ivp/_zeros.h @@ -0,0 +1,40 @@ +/* Written by Charles Harris charles.harris@sdl.usu.edu */ + +/* Modified to not depend on Python everywhere by Travis Oliphant. + */ + +#ifndef ZEROS_H +#define ZEROS_H + +typedef struct { + int funcalls; + int iterations; + int error_num; +} scipy_zeros_info; + + +/* Must agree with _ECONVERGED, _ESIGNERR, _ECONVERR in zeros.py */ +#define CONVERGED 0 +#define SIGNERR -1 +#define CONVERR -2 +#define EVALUEERR -3 +#define INPROGRESS 1 + +typedef double (*callback_type)(double, void*); +typedef double (*solver_type)(callback_type, double, double, double, double, + int, void *, scipy_zeros_info*); + +extern double bisect(callback_type f, double xa, double xb, double xtol, + double rtol, int iter, void *func_data_param, + scipy_zeros_info *solver_stats); +extern double ridder(callback_type f, double xa, double xb, double xtol, + double rtol, int iter, void *func_data_param, + scipy_zeros_info *solver_stats); +extern double brenth(callback_type f, double xa, double xb, double xtol, + double rtol, int iter, void *func_data_param, + scipy_zeros_info *solver_stats); +extern double brentq(callback_type f, double xa, double xb, double xtol, + double rtol, int iter, void *func_data_param, + scipy_zeros_info *solver_stats); + +#endif From 10b49cacf1113c787c552c1e4925acda8b12b3f3 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 10 Jan 2024 10:35:59 +0100 Subject: [PATCH 082/346] solver isolated, works --- src/hapsira/core/math/ivp/_brentq.py | 4 +- src/hapsira/core/math/ivp/_solver.py | 151 +++++++++++++++++++++++++++ src/hapsira/core/math/ivp/_zeros.c | 49 +++++++++ 3 files changed, 202 insertions(+), 2 deletions(-) create mode 100644 src/hapsira/core/math/ivp/_solver.py diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index 50993f4db..64d6f64e6 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -3,7 +3,7 @@ import numpy as np -# _zeros, OptimizeResult +from ._solver import brentq_sf _ECONVERGED = 0 @@ -304,5 +304,5 @@ def brentq( if rtol < _rtol: raise ValueError(f"rtol too small ({rtol:g} < {_rtol:g})") f = _wrap_nan_raise(f) - r = _zeros._brentq(f, a, b, xtol, rtol, maxiter, args, full_output, disp) + r = brentq_sf(f, a, b, xtol, rtol, maxiter, args, full_output, disp) return results_c(full_output, r) diff --git a/src/hapsira/core/math/ivp/_solver.py b/src/hapsira/core/math/ivp/_solver.py new file mode 100644 index 000000000..011d4e611 --- /dev/null +++ b/src/hapsira/core/math/ivp/_solver.py @@ -0,0 +1,151 @@ + +from numba import njit, jit + +from math import fabs + + +CONVERGED = 0 +SIGNERR = -1 +CONVERR = -2 +# EVALUEERR = -3 +INPROGRESS = 1 + + +@njit +def _min_hf(a, b): + return a if a < b else b + + +@njit +def _signbit_hf(a): + return a < 0 + + +@jit +def _brentq_hf( + func, # callback_type + xa, # double + xb, # double + xtol, # double + rtol, # double + iter_, # int + xargs, # void (*) +): + + xpre, xcur = xa, xb + xblk = 0. + fpre, fcur, fblk = 0., 0., 0. + spre, scur = 0., 0. + + iterations = 0 + + fpre = func(xpre, *xargs) + fcur = func(xcur, *xargs) + funcalls = 2 + if fpre == 0: + return xpre, funcalls, iterations, CONVERGED + if fcur == 0: + return xcur, funcalls, iterations, CONVERGED + if _signbit_hf(fpre) == _signbit_hf(fcur): + return 0., funcalls, iterations, SIGNERR + + iterations = 0 + for _ in range(0, iter_): + iterations += 1 + if fpre != 0 and fcur != 0 and _signbit_hf(fpre) != _signbit_hf(fcur): + xblk = xpre + fblk = fpre + scur = xcur - xpre + spre = scur + if fabs(fblk) < fabs(fcur): + xpre = xcur + xcur = xblk + xblk = xpre + + fpre = fcur + fcur = fblk + fblk = fpre + + delta = (xtol + rtol * fabs(xcur)) / 2 + sbis = (xblk - xcur) / 2 + if fcur == 0 or fabs(sbis) < delta: + return xcur, funcalls, iterations, CONVERGED + + if fabs(spre) > delta and fabs(fcur) < fabs(fpre): + if xpre == xblk: + stry = -fcur * (xcur - xpre) / (fcur - fpre) + else: + dpre = (fpre - fcur) / (xpre - xcur) + dblk = (fblk - fcur) / (xblk - xcur) + stry = -fcur * (fblk * dblk - fpre * dpre) / (dblk * dpre * (fblk - fpre)) + if 2 * fabs(stry) < _min_hf(fabs(spre), 3 * fabs(sbis) - delta): + spre = scur + scur = stry + else: + spre = sbis + scur = sbis + else: + spre = sbis + scur = sbis + + xpre = xcur + fpre = fcur + if fabs(scur) > delta: + xcur += scur + else: + xcur += delta if sbis > 0 else -delta + + fcur = func(xcur, *xargs) + funcalls += 1 + + return xcur, funcalls, iterations, CONVERR + + +@jit +def brentq_sf( + func, # func + a, # double + b, # double + xtol, # double + rtol, # double + iter_, # int + xargs, # args tuple for func + fulloutput, # int + disp, # int +): + + if xtol < 0: + raise ValueError("xtol must be >= 0") + if iter_ < 0: + raise ValueError("maxiter should be > 0") + + zero, funcalls, iterations, error_num = _brentq_hf( + func, + a, + b, + xtol, + rtol, + iter_, + xargs, + ) + + flag = 0 + if error_num != CONVERGED: + if error_num == SIGNERR: + raise ValueError("f(a) and f(b) must have different signs") + if error_num == CONVERR: + if disp: + raise RuntimeError("Failed to converge after %d iterations." % iterations) + flag = CONVERR + else: + flag = CONVERGED + + if fulloutput: + return ( + zero, # double + funcalls, # int + iterations, # int + flag, # int + ) + else: + return zero # double diff --git a/src/hapsira/core/math/ivp/_zeros.c b/src/hapsira/core/math/ivp/_zeros.c index 0516ffab9..f8f13bda2 100644 --- a/src/hapsira/core/math/ivp/_zeros.c +++ b/src/hapsira/core/math/ivp/_zeros.c @@ -1,3 +1,52 @@ + +typedef struct { + PyObject *function; + PyObject *xargs; + jmp_buf env; +} scipy_zeros_parameters; + + + +static double +scipy_zeros_functions_func(double x, void *params) +{ + scipy_zeros_parameters *myparams = params; + PyObject *args, *xargs, *item, *f, *retval=NULL; + Py_ssize_t i, len; + double val; + + xargs = myparams->xargs; + /* Need to create a new 'args' tuple on each call in case 'f' is + stateful and keeps references to it (e.g. functools.lru_cache) */ + len = PyTuple_Size(xargs); + /* Make room for the double as first argument */ + args = PyArgs(New)(len + 1); + if (args == NULL) { + PyErr_SetString(PyExc_RuntimeError, "Failed to allocate arguments"); + longjmp(myparams->env, 1); + } + PyArgs(SET_ITEM)(args, 0, Py_BuildValue("d", x)); + for (i = 0; i < len; i++) { + item = PyTuple_GetItem(xargs, i); + if (item == NULL) { + Py_DECREF(args); + longjmp(myparams->env, 1); + } + Py_INCREF(item); + PyArgs(SET_ITEM)(args, i+1, item); + } + + f = myparams->function; + retval = PyObject_CallObject(f,args); + Py_DECREF(args); + if (retval == NULL) { + longjmp(myparams->env, 1); + } + val = PyFloat_AsDouble(retval); + Py_XDECREF(retval); + return val; +} + /* * Helper function that calls a Python function with extended arguments */ From 91e68f20e0a6f9b051ee7debf81ee74fefdfd515 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 10 Jan 2024 10:36:11 +0100 Subject: [PATCH 083/346] solver isolated, works --- src/hapsira/core/math/ivp/_solver.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solver.py b/src/hapsira/core/math/ivp/_solver.py index 011d4e611..7d3375a57 100644 --- a/src/hapsira/core/math/ivp/_solver.py +++ b/src/hapsira/core/math/ivp/_solver.py @@ -1,4 +1,3 @@ - from numba import njit, jit from math import fabs @@ -31,11 +30,10 @@ def _brentq_hf( iter_, # int xargs, # void (*) ): - xpre, xcur = xa, xb - xblk = 0. - fpre, fcur, fblk = 0., 0., 0. - spre, scur = 0., 0. + xblk = 0.0 + fpre, fcur, fblk = 0.0, 0.0, 0.0 + spre, scur = 0.0, 0.0 iterations = 0 @@ -47,7 +45,7 @@ def _brentq_hf( if fcur == 0: return xcur, funcalls, iterations, CONVERGED if _signbit_hf(fpre) == _signbit_hf(fcur): - return 0., funcalls, iterations, SIGNERR + return 0.0, funcalls, iterations, SIGNERR iterations = 0 for _ in range(0, iter_): @@ -77,7 +75,9 @@ def _brentq_hf( else: dpre = (fpre - fcur) / (xpre - xcur) dblk = (fblk - fcur) / (xblk - xcur) - stry = -fcur * (fblk * dblk - fpre * dpre) / (dblk * dpre * (fblk - fpre)) + stry = ( + -fcur * (fblk * dblk - fpre * dpre) / (dblk * dpre * (fblk - fpre)) + ) if 2 * fabs(stry) < _min_hf(fabs(spre), 3 * fabs(sbis) - delta): spre = scur scur = stry @@ -113,7 +113,6 @@ def brentq_sf( fulloutput, # int disp, # int ): - if xtol < 0: raise ValueError("xtol must be >= 0") if iter_ < 0: @@ -135,7 +134,9 @@ def brentq_sf( raise ValueError("f(a) and f(b) must have different signs") if error_num == CONVERR: if disp: - raise RuntimeError("Failed to converge after %d iterations." % iterations) + raise RuntimeError( + "Failed to converge after %d iterations." % iterations + ) flag = CONVERR else: flag = CONVERGED From 5a7f521e0e534e932f021abf281f77dd103e5437 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 10 Jan 2024 10:38:31 +0100 Subject: [PATCH 084/346] rm c/h --- src/hapsira/core/math/ivp/_brentq.c | 130 -------------------- src/hapsira/core/math/ivp/_zeros.c | 181 ---------------------------- src/hapsira/core/math/ivp/_zeros.h | 40 ------ 3 files changed, 351 deletions(-) delete mode 100644 src/hapsira/core/math/ivp/_brentq.c delete mode 100644 src/hapsira/core/math/ivp/_zeros.c delete mode 100644 src/hapsira/core/math/ivp/_zeros.h diff --git a/src/hapsira/core/math/ivp/_brentq.c b/src/hapsira/core/math/ivp/_brentq.c deleted file mode 100644 index 1d8df1eb8..000000000 --- a/src/hapsira/core/math/ivp/_brentq.c +++ /dev/null @@ -1,130 +0,0 @@ -/* Written by Charles Harris charles.harris@sdl.usu.edu */ - -#include -#include "zeros.h" - -#define MIN(a, b) ((a) < (b) ? (a) : (b)) - -/* - At the top of the loop the situation is the following: - - 1. the root is bracketed between xa and xb - 2. xa is the most recent estimate - 3. xp is the previous estimate - 4. |fp| < |fb| - - The order of xa and xp doesn't matter, but assume xp < xb. Then xa lies to - the right of xp and the assumption is that xa is increasing towards the root. - In this situation we will attempt quadratic extrapolation as long as the - condition - - * |fa| < |fp| < |fb| - - is satisfied. That is, the function value is decreasing as we go along. - Note the 4 above implies that the right inequlity already holds. - - The first check is that xa is still to the left of the root. If not, xb is - replaced by xp and the interval reverses, with xb < xa. In this situation - we will try linear interpolation. That this has happened is signaled by the - equality xb == xp; - - The second check is that |fa| < |fb|. If this is not the case, we swap - xa and xb and resort to bisection. - -*/ - -double -brentq(callback_type f, double xa, double xb, double xtol, double rtol, - int iter, void *func_data_param, scipy_zeros_info *solver_stats) -{ - double xpre = xa, xcur = xb; - double xblk = 0., fpre, fcur, fblk = 0., spre = 0., scur = 0., sbis; - /* the tolerance is 2*delta */ - double delta; - double stry, dpre, dblk; - int i; - solver_stats->error_num = INPROGRESS; - - fpre = (*f)(xpre, func_data_param); - fcur = (*f)(xcur, func_data_param); - solver_stats->funcalls = 2; - if (fpre == 0) { - solver_stats->error_num = CONVERGED; - return xpre; - } - if (fcur == 0) { - solver_stats->error_num = CONVERGED; - return xcur; - } - if (signbit(fpre)==signbit(fcur)) { - solver_stats->error_num = SIGNERR; - return 0.; - } - solver_stats->iterations = 0; - for (i = 0; i < iter; i++) { - solver_stats->iterations++; - if (fpre != 0 && fcur != 0 && - (signbit(fpre) != signbit(fcur))) { - xblk = xpre; - fblk = fpre; - spre = scur = xcur - xpre; - } - if (fabs(fblk) < fabs(fcur)) { - xpre = xcur; - xcur = xblk; - xblk = xpre; - - fpre = fcur; - fcur = fblk; - fblk = fpre; - } - - delta = (xtol + rtol*fabs(xcur))/2; - sbis = (xblk - xcur)/2; - if (fcur == 0 || fabs(sbis) < delta) { - solver_stats->error_num = CONVERGED; - return xcur; - } - - if (fabs(spre) > delta && fabs(fcur) < fabs(fpre)) { - if (xpre == xblk) { - /* interpolate */ - stry = -fcur*(xcur - xpre)/(fcur - fpre); - } - else { - /* extrapolate */ - dpre = (fpre - fcur)/(xpre - xcur); - dblk = (fblk - fcur)/(xblk - xcur); - stry = -fcur*(fblk*dblk - fpre*dpre) - /(dblk*dpre*(fblk - fpre)); - } - if (2*fabs(stry) < MIN(fabs(spre), 3*fabs(sbis) - delta)) { - /* good short step */ - spre = scur; - scur = stry; - } else { - /* bisect */ - spre = sbis; - scur = sbis; - } - } - else { - /* bisect */ - spre = sbis; - scur = sbis; - } - - xpre = xcur; fpre = fcur; - if (fabs(scur) > delta) { - xcur += scur; - } - else { - xcur += (sbis > 0 ? delta : -delta); - } - - fcur = (*f)(xcur, func_data_param); - solver_stats->funcalls++; - } - solver_stats->error_num = CONVERR; - return xcur; -} diff --git a/src/hapsira/core/math/ivp/_zeros.c b/src/hapsira/core/math/ivp/_zeros.c deleted file mode 100644 index f8f13bda2..000000000 --- a/src/hapsira/core/math/ivp/_zeros.c +++ /dev/null @@ -1,181 +0,0 @@ - -typedef struct { - PyObject *function; - PyObject *xargs; - jmp_buf env; -} scipy_zeros_parameters; - - - -static double -scipy_zeros_functions_func(double x, void *params) -{ - scipy_zeros_parameters *myparams = params; - PyObject *args, *xargs, *item, *f, *retval=NULL; - Py_ssize_t i, len; - double val; - - xargs = myparams->xargs; - /* Need to create a new 'args' tuple on each call in case 'f' is - stateful and keeps references to it (e.g. functools.lru_cache) */ - len = PyTuple_Size(xargs); - /* Make room for the double as first argument */ - args = PyArgs(New)(len + 1); - if (args == NULL) { - PyErr_SetString(PyExc_RuntimeError, "Failed to allocate arguments"); - longjmp(myparams->env, 1); - } - PyArgs(SET_ITEM)(args, 0, Py_BuildValue("d", x)); - for (i = 0; i < len; i++) { - item = PyTuple_GetItem(xargs, i); - if (item == NULL) { - Py_DECREF(args); - longjmp(myparams->env, 1); - } - Py_INCREF(item); - PyArgs(SET_ITEM)(args, i+1, item); - } - - f = myparams->function; - retval = PyObject_CallObject(f,args); - Py_DECREF(args); - if (retval == NULL) { - longjmp(myparams->env, 1); - } - val = PyFloat_AsDouble(retval); - Py_XDECREF(retval); - return val; -} - -/* - * Helper function that calls a Python function with extended arguments - */ - -static PyObject * -call_solver(solver_type solver, PyObject *self, PyObject *args) -{ - double a, b, xtol, rtol, zero; - int iter, fulloutput, disp=1, flag=0; - scipy_zeros_parameters params; - scipy_zeros_info solver_stats; - PyObject *f, *xargs; - - if (!PyArg_ParseTuple(args, "OddddiOi|i", - &f, &a, &b, &xtol, &rtol, &iter, &xargs, &fulloutput, &disp)) { - PyErr_SetString(PyExc_RuntimeError, "Unable to parse arguments"); - return NULL; - } - if (xtol < 0) { - PyErr_SetString(PyExc_ValueError, "xtol must be >= 0"); - return NULL; - } - if (iter < 0) { - PyErr_SetString(PyExc_ValueError, "maxiter should be > 0"); - return NULL; - } - - params.function = f; - params.xargs = xargs; - - if (!setjmp(params.env)) { - /* direct return */ - solver_stats.error_num = 0; - zero = solver(scipy_zeros_functions_func, a, b, xtol, rtol, - iter, (void*)¶ms, &solver_stats); - } else { - /* error return from Python function */ - return NULL; - } - - if (solver_stats.error_num != CONVERGED) { - if (solver_stats.error_num == SIGNERR) { - PyErr_SetString(PyExc_ValueError, - "f(a) and f(b) must have different signs"); - return NULL; - } - if (solver_stats.error_num == CONVERR) { - if (disp) { - char msg[100]; - PyOS_snprintf(msg, sizeof(msg), - "Failed to converge after %d iterations.", - solver_stats.iterations); - PyErr_SetString(PyExc_RuntimeError, msg); - return NULL; - } - flag = CONVERR; - } - } - else { - flag = CONVERGED; - } - if (fulloutput) { - return Py_BuildValue("diii", - zero, solver_stats.funcalls, solver_stats.iterations, flag); - } - else { - return Py_BuildValue("d", zero); - } -} - -/* - * These routines interface with the solvers through call_solver - */ - -static PyObject * -_bisect(PyObject *self, PyObject *args) -{ - return call_solver(bisect,self,args); -} - -static PyObject * -_ridder(PyObject *self, PyObject *args) -{ - return call_solver(ridder,self,args); -} - -static PyObject * -_brenth(PyObject *self, PyObject *args) -{ - return call_solver(brenth,self,args); -} - -static PyObject * -_brentq(PyObject *self, PyObject *args) -{ - return call_solver(brentq,self,args); -} - -/* - * Standard Python module interface - */ - -static PyMethodDef -Zerosmethods[] = { - {"_bisect", _bisect, METH_VARARGS, "a"}, - {"_ridder", _ridder, METH_VARARGS, "a"}, - {"_brenth", _brenth, METH_VARARGS, "a"}, - {"_brentq", _brentq, METH_VARARGS, "a"}, - {NULL, NULL} -}; - -static struct PyModuleDef moduledef = { - PyModuleDef_HEAD_INIT, - "_zeros", - NULL, - -1, - Zerosmethods, - NULL, - NULL, - NULL, - NULL -}; - -PyMODINIT_FUNC -PyInit__zeros(void) -{ - PyObject *m; - - m = PyModule_Create(&moduledef); - - return m; -} diff --git a/src/hapsira/core/math/ivp/_zeros.h b/src/hapsira/core/math/ivp/_zeros.h deleted file mode 100644 index 14afcc13f..000000000 --- a/src/hapsira/core/math/ivp/_zeros.h +++ /dev/null @@ -1,40 +0,0 @@ -/* Written by Charles Harris charles.harris@sdl.usu.edu */ - -/* Modified to not depend on Python everywhere by Travis Oliphant. - */ - -#ifndef ZEROS_H -#define ZEROS_H - -typedef struct { - int funcalls; - int iterations; - int error_num; -} scipy_zeros_info; - - -/* Must agree with _ECONVERGED, _ESIGNERR, _ECONVERR in zeros.py */ -#define CONVERGED 0 -#define SIGNERR -1 -#define CONVERR -2 -#define EVALUEERR -3 -#define INPROGRESS 1 - -typedef double (*callback_type)(double, void*); -typedef double (*solver_type)(callback_type, double, double, double, double, - int, void *, scipy_zeros_info*); - -extern double bisect(callback_type f, double xa, double xb, double xtol, - double rtol, int iter, void *func_data_param, - scipy_zeros_info *solver_stats); -extern double ridder(callback_type f, double xa, double xb, double xtol, - double rtol, int iter, void *func_data_param, - scipy_zeros_info *solver_stats); -extern double brenth(callback_type f, double xa, double xb, double xtol, - double rtol, int iter, void *func_data_param, - scipy_zeros_info *solver_stats); -extern double brentq(callback_type f, double xa, double xb, double xtol, - double rtol, int iter, void *func_data_param, - scipy_zeros_info *solver_stats); - -#endif From 3b26405046ec3a9878bc28e1c186d85fbdffc13f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 10 Jan 2024 11:20:20 +0100 Subject: [PATCH 085/346] expose brentq --- src/hapsira/core/math/ivp/__init__.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/ivp/__init__.py b/src/hapsira/core/math/ivp/__init__.py index da2192fdc..4c1c497c1 100644 --- a/src/hapsira/core/math/ivp/__init__.py +++ b/src/hapsira/core/math/ivp/__init__.py @@ -1,5 +1,4 @@ from ._api import solve_ivp +from ._brentq import brentq -__all__ = [ - "solve_ivp", -] +__all__ = ["solve_ivp", "brentq"] From 7fb01f07e566da4a7ad4e202169c78c1e19f53fc Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 10 Jan 2024 11:20:56 +0100 Subject: [PATCH 086/346] rm links to c impl --- src/hapsira/core/math/ivp/_api.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index b01bb8e30..9917e24af 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -119,8 +119,6 @@ def solve_event_equation(event, sol, t_old, t): """ return brentq(lambda t: event(t, sol(t)), t_old, t, xtol=4 * EPS, rtol=4 * EPS) - # https://github.com/scipy/scipy/blob/ea4d1f1330950bee74396e427fe6330424907621/scipy/optimize/_zeros_py.py#L682 - # https://github.com/scipy/scipy/blob/ea4d1f1330950bee74396e427fe6330424907621/scipy/optimize/Zeros/brentq.c#L37 def handle_events(sol, events, active_events, is_terminal, t_old, t): From b8273f7216f9b9be2501e66c362317c39da382dc Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 10 Jan 2024 11:21:29 +0100 Subject: [PATCH 087/346] use internal brentq --- src/hapsira/core/math/optimize.py | 5 ----- 1 file changed, 5 deletions(-) delete mode 100644 src/hapsira/core/math/optimize.py diff --git a/src/hapsira/core/math/optimize.py b/src/hapsira/core/math/optimize.py deleted file mode 100644 index 4033e6080..000000000 --- a/src/hapsira/core/math/optimize.py +++ /dev/null @@ -1,5 +0,0 @@ -from scipy.optimize import brentq - -__all__ = [ - "brentq", -] From 227456a14a927e65184edf4075636c0e84d9f812 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 10 Jan 2024 11:21:44 +0100 Subject: [PATCH 088/346] use internal brentq --- src/hapsira/threebody/restricted.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/threebody/restricted.py b/src/hapsira/threebody/restricted.py index 58b37ba67..67ffdf66d 100644 --- a/src/hapsira/threebody/restricted.py +++ b/src/hapsira/threebody/restricted.py @@ -7,7 +7,7 @@ from astropy import units as u import numpy as np -from hapsira.core.math.optimize import brentq +from hapsira.core.math.ivp import brentq from hapsira.util import norm From 99d1e59a87b5dd6e596c2eb36a8e33c7b290b5f8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 10 Jan 2024 11:22:21 +0100 Subject: [PATCH 089/346] cleanup --- src/hapsira/core/math/ivp/_brentq.py | 127 +-------------------------- src/hapsira/core/math/ivp/_solver.py | 30 ++----- 2 files changed, 8 insertions(+), 149 deletions(-) diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index 64d6f64e6..75b8ffc0f 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -27,116 +27,6 @@ } -class OptimizeResult(dict): - """Represents the optimization result. - - Attributes - ---------- - x : ndarray - The solution of the optimization. - success : bool - Whether or not the optimizer exited successfully. - status : int - Termination status of the optimizer. Its value depends on the - underlying solver. Refer to `message` for details. - message : str - Description of the cause of the termination. - fun, jac, hess: ndarray - Values of objective function, its Jacobian and its Hessian (if - available). The Hessians may be approximations, see the documentation - of the function in question. - hess_inv : object - Inverse of the objective function's Hessian; may be an approximation. - Not available for all solvers. The type of this attribute may be - either np.ndarray or scipy.sparse.linalg.LinearOperator. - nfev, njev, nhev : int - Number of evaluations of the objective functions and of its - Jacobian and Hessian. - nit : int - Number of iterations performed by the optimizer. - maxcv : float - The maximum constraint violation. - - Notes - ----- - Depending on the specific solver being used, `OptimizeResult` may - not have all attributes listed here, and they may have additional - attributes not listed here. Since this class is essentially a - subclass of dict with attribute accessors, one can see which - attributes are available using the `OptimizeResult.keys` method. - """ - - def __getattr__(self, name): - try: - return self[name] - except KeyError as e: - raise AttributeError(name) from e - - __setattr__ = dict.__setitem__ - __delattr__ = dict.__delitem__ - - # def __repr__(self): - # order_keys = ['message', 'success', 'status', 'fun', 'funl', 'x', 'xl', - # 'col_ind', 'nit', 'lower', 'upper', 'eqlin', 'ineqlin', - # 'converged', 'flag', 'function_calls', 'iterations', - # 'root'] - # # 'slack', 'con' are redundant with residuals - # # 'crossover_nit' is probably not interesting to most users - # omit_keys = {'slack', 'con', 'crossover_nit'} - - # def key(item): - # try: - # return order_keys.index(item[0].lower()) - # except ValueError: # item not in list - # return np.inf - - # def omit_redundant(items): - # for item in items: - # if item[0] in omit_keys: - # continue - # yield item - - # def item_sorter(d): - # return sorted(omit_redundant(d.items()), key=key) - - # if self.keys(): - # return _dict_formatter(self, sorter=item_sorter) - # else: - # return self.__class__.__name__ + "()" - - def __dir__(self): - return list(self.keys()) - - -class RootResults(OptimizeResult): - """Represents the root finding result. - - Attributes - ---------- - root : float - Estimated root location. - iterations : int - Number of iterations needed to find the root. - function_calls : int - Number of times the function was called. - converged : bool - True if the routine converged. - flag : str - Description of the cause of termination. - - """ - - def __init__(self, root, iterations, function_calls, flag): - self.root = root - self.iterations = iterations - self.function_calls = function_calls - self.converged = flag == _ECONVERGED - if flag in flag_map: - self.flag = flag_map[flag] - else: - self.flag = flag - - def _wrap_nan_raise(f): def f_raise(x, *args): fx = f(x, *args) @@ -153,17 +43,6 @@ def f_raise(x, *args): return f_raise -def results_c(full_output, r): - if full_output: - x, funcalls, iterations, flag = r - results = RootResults( - root=x, iterations=iterations, function_calls=funcalls, flag=flag - ) - return x, results - else: - return r - - _iter = 100 _xtol = 2e-12 _rtol = 4 * np.finfo(float).eps @@ -177,8 +56,6 @@ def brentq( xtol=_xtol, rtol=_rtol, maxiter=_iter, - full_output=False, - disp=True, ): """ Find a root of a function in a bracketing interval using Brent's method. @@ -304,5 +181,5 @@ def brentq( if rtol < _rtol: raise ValueError(f"rtol too small ({rtol:g} < {_rtol:g})") f = _wrap_nan_raise(f) - r = brentq_sf(f, a, b, xtol, rtol, maxiter, args, full_output, disp) - return results_c(full_output, r) + r = brentq_sf(f, a, b, xtol, rtol, maxiter, args) + return r diff --git a/src/hapsira/core/math/ivp/_solver.py b/src/hapsira/core/math/ivp/_solver.py index 7d3375a57..6156ea113 100644 --- a/src/hapsira/core/math/ivp/_solver.py +++ b/src/hapsira/core/math/ivp/_solver.py @@ -110,8 +110,6 @@ def brentq_sf( rtol, # double iter_, # int xargs, # args tuple for func - fulloutput, # int - disp, # int ): if xtol < 0: raise ValueError("xtol must be >= 0") @@ -128,25 +126,9 @@ def brentq_sf( xargs, ) - flag = 0 - if error_num != CONVERGED: - if error_num == SIGNERR: - raise ValueError("f(a) and f(b) must have different signs") - if error_num == CONVERR: - if disp: - raise RuntimeError( - "Failed to converge after %d iterations." % iterations - ) - flag = CONVERR - else: - flag = CONVERGED - - if fulloutput: - return ( - zero, # double - funcalls, # int - iterations, # int - flag, # int - ) - else: - return zero # double + if error_num == SIGNERR: + raise ValueError("f(a) and f(b) must have different signs") + if error_num == CONVERR: + raise RuntimeError("Failed to converge after %d iterations." % iterations) + + return zero # double From 292908b60432ab352d3f625e74534015887e0601 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 10 Jan 2024 14:53:43 +0100 Subject: [PATCH 090/346] drop arb args in solver; cleanup of unused consts --- src/hapsira/core/math/ivp/_brentq.py | 33 +++------------------------- src/hapsira/core/math/ivp/_solver.py | 9 +++----- 2 files changed, 6 insertions(+), 36 deletions(-) diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index 75b8ffc0f..3126bd593 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -6,30 +6,9 @@ from ._solver import brentq_sf -_ECONVERGED = 0 -_ESIGNERR = -1 -_ECONVERR = -2 -_EVALUEERR = -3 -_EINPROGRESS = 1 - -CONVERGED = "converged" -SIGNERR = "sign error" -CONVERR = "convergence error" -VALUEERR = "value error" -INPROGRESS = "No error" - -flag_map = { - _ECONVERGED: CONVERGED, - _ESIGNERR: SIGNERR, - _ECONVERR: CONVERR, - _EVALUEERR: VALUEERR, - _EINPROGRESS: INPROGRESS, -} - - def _wrap_nan_raise(f): - def f_raise(x, *args): - fx = f(x, *args) + def f_raise(x): + fx = f(x) f_raise._function_calls += 1 if np.isnan(fx): msg = f"The function value at x={x} is NaN; " "solver cannot continue." @@ -52,7 +31,6 @@ def brentq( f, a, b, - args=(), xtol=_xtol, rtol=_rtol, maxiter=_iter, @@ -102,9 +80,6 @@ def brentq( maxiter : int, optional If convergence is not achieved in `maxiter` iterations, an error is raised. Must be >= 0. - args : tuple, optional - Containing extra arguments for the function `f`. - `f` is called by ``apply(f, (x)+args)``. full_output : bool, optional If `full_output` is False, the root is returned. If `full_output` is True, the return value is ``(x, r)``, where `x` is the root, and `r` is @@ -173,13 +148,11 @@ def brentq( >>> root 1.0 """ - if not isinstance(args, tuple): - args = (args,) maxiter = operator.index(maxiter) if xtol <= 0: raise ValueError("xtol too small (%g <= 0)" % xtol) if rtol < _rtol: raise ValueError(f"rtol too small ({rtol:g} < {_rtol:g})") f = _wrap_nan_raise(f) - r = brentq_sf(f, a, b, xtol, rtol, maxiter, args) + r = brentq_sf(f, a, b, xtol, rtol, maxiter) return r diff --git a/src/hapsira/core/math/ivp/_solver.py b/src/hapsira/core/math/ivp/_solver.py index 6156ea113..ce68a76e7 100644 --- a/src/hapsira/core/math/ivp/_solver.py +++ b/src/hapsira/core/math/ivp/_solver.py @@ -28,7 +28,6 @@ def _brentq_hf( xtol, # double rtol, # double iter_, # int - xargs, # void (*) ): xpre, xcur = xa, xb xblk = 0.0 @@ -37,8 +36,8 @@ def _brentq_hf( iterations = 0 - fpre = func(xpre, *xargs) - fcur = func(xcur, *xargs) + fpre = func(xpre) + fcur = func(xcur) funcalls = 2 if fpre == 0: return xpre, funcalls, iterations, CONVERGED @@ -95,7 +94,7 @@ def _brentq_hf( else: xcur += delta if sbis > 0 else -delta - fcur = func(xcur, *xargs) + fcur = func(xcur) funcalls += 1 return xcur, funcalls, iterations, CONVERR @@ -109,7 +108,6 @@ def brentq_sf( xtol, # double rtol, # double iter_, # int - xargs, # args tuple for func ): if xtol < 0: raise ValueError("xtol must be >= 0") @@ -123,7 +121,6 @@ def brentq_sf( xtol, rtol, iter_, - xargs, ) if error_num == SIGNERR: From 2b7cf898088603b808d923e61d882b39da4bd2a3 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 10 Jan 2024 15:03:36 +0100 Subject: [PATCH 091/346] all brentq in one place --- src/hapsira/core/math/ivp/_brentq.py | 262 ++++++++++++++------------- src/hapsira/core/math/ivp/_solver.py | 131 -------------- 2 files changed, 140 insertions(+), 253 deletions(-) delete mode 100644 src/hapsira/core/math/ivp/_solver.py diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index 3126bd593..b9d6993e6 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -1,9 +1,141 @@ +from math import fabs + import operator +from numba import njit, jit import numpy as np -from ._solver import brentq_sf +CONVERGED = 0 +SIGNERR = -1 +CONVERR = -2 +# EVALUEERR = -3 +INPROGRESS = 1 + +BRENTQ_ITER = 100 +BRENTQ_XTOL = 2e-12 +BRENTQ_RTOL = 4 * np.finfo(float).eps + + +@njit +def _min_hf(a, b): + return a if a < b else b + + +@njit +def _signbit_hf(a): + return a < 0 + + +@jit +def _brentq_hf( + func, # callback_type + xa, # double + xb, # double + xtol, # double + rtol, # double + iter_, # int +): + xpre, xcur = xa, xb + xblk = 0.0 + fpre, fcur, fblk = 0.0, 0.0, 0.0 + spre, scur = 0.0, 0.0 + + iterations = 0 + + fpre = func(xpre) + fcur = func(xcur) + funcalls = 2 + if fpre == 0: + return xpre, funcalls, iterations, CONVERGED + if fcur == 0: + return xcur, funcalls, iterations, CONVERGED + if _signbit_hf(fpre) == _signbit_hf(fcur): + return 0.0, funcalls, iterations, SIGNERR + + iterations = 0 + for _ in range(0, iter_): + iterations += 1 + if fpre != 0 and fcur != 0 and _signbit_hf(fpre) != _signbit_hf(fcur): + xblk = xpre + fblk = fpre + scur = xcur - xpre + spre = scur + if fabs(fblk) < fabs(fcur): + xpre = xcur + xcur = xblk + xblk = xpre + + fpre = fcur + fcur = fblk + fblk = fpre + + delta = (xtol + rtol * fabs(xcur)) / 2 + sbis = (xblk - xcur) / 2 + if fcur == 0 or fabs(sbis) < delta: + return xcur, funcalls, iterations, CONVERGED + + if fabs(spre) > delta and fabs(fcur) < fabs(fpre): + if xpre == xblk: + stry = -fcur * (xcur - xpre) / (fcur - fpre) + else: + dpre = (fpre - fcur) / (xpre - xcur) + dblk = (fblk - fcur) / (xblk - xcur) + stry = ( + -fcur * (fblk * dblk - fpre * dpre) / (dblk * dpre * (fblk - fpre)) + ) + if 2 * fabs(stry) < _min_hf(fabs(spre), 3 * fabs(sbis) - delta): + spre = scur + scur = stry + else: + spre = sbis + scur = sbis + else: + spre = sbis + scur = sbis + + xpre = xcur + fpre = fcur + if fabs(scur) > delta: + xcur += scur + else: + xcur += delta if sbis > 0 else -delta + + fcur = func(xcur) + funcalls += 1 + + return xcur, funcalls, iterations, CONVERR + + +@jit +def brentq_sf( + func, # func + a, # double + b, # double + xtol, # double + rtol, # double + iter_, # int +): + if xtol < 0: + raise ValueError("xtol must be >= 0") + if iter_ < 0: + raise ValueError("maxiter should be > 0") + + zero, funcalls, iterations, error_num = _brentq_hf( + func, + a, + b, + xtol, + rtol, + iter_, + ) + + if error_num == SIGNERR: + raise ValueError("f(a) and f(b) must have different signs") + if error_num == CONVERR: + raise RuntimeError("Failed to converge after %d iterations." % iterations) + + return zero # double def _wrap_nan_raise(f): @@ -22,137 +154,23 @@ def f_raise(x): return f_raise -_iter = 100 -_xtol = 2e-12 -_rtol = 4 * np.finfo(float).eps - - def brentq( f, a, b, - xtol=_xtol, - rtol=_rtol, - maxiter=_iter, + xtol=BRENTQ_XTOL, + rtol=BRENTQ_RTOL, + maxiter=BRENTQ_ITER, ): """ - Find a root of a function in a bracketing interval using Brent's method. - - Uses the classic Brent's method to find a root of the function `f` on - the sign changing interval [a , b]. Generally considered the best of the - rootfinding routines here. It is a safe version of the secant method that - uses inverse quadratic extrapolation. Brent's method combines root - bracketing, interval bisection, and inverse quadratic interpolation. It is - sometimes known as the van Wijngaarden-Dekker-Brent method. Brent (1973) - claims convergence is guaranteed for functions computable within [a,b]. - - [Brent1973]_ provides the classic description of the algorithm. Another - description can be found in a recent edition of Numerical Recipes, including - [PressEtal1992]_. A third description is at - http://mathworld.wolfram.com/BrentsMethod.html. It should be easy to - understand the algorithm just by reading our code. Our code diverges a bit - from standard presentations: we choose a different formula for the - extrapolation step. - - Parameters - ---------- - f : function - Python function returning a number. The function :math:`f` - must be continuous, and :math:`f(a)` and :math:`f(b)` must - have opposite signs. - a : scalar - One end of the bracketing interval :math:`[a, b]`. - b : scalar - The other end of the bracketing interval :math:`[a, b]`. - xtol : number, optional - The computed root ``x0`` will satisfy ``np.allclose(x, x0, - atol=xtol, rtol=rtol)``, where ``x`` is the exact root. The - parameter must be positive. For nice functions, Brent's - method will often satisfy the above condition with ``xtol/2`` - and ``rtol/2``. [Brent1973]_ - rtol : number, optional - The computed root ``x0`` will satisfy ``np.allclose(x, x0, - atol=xtol, rtol=rtol)``, where ``x`` is the exact root. The - parameter cannot be smaller than its default value of - ``4*np.finfo(float).eps``. For nice functions, Brent's - method will often satisfy the above condition with ``xtol/2`` - and ``rtol/2``. [Brent1973]_ - maxiter : int, optional - If convergence is not achieved in `maxiter` iterations, an error is - raised. Must be >= 0. - full_output : bool, optional - If `full_output` is False, the root is returned. If `full_output` is - True, the return value is ``(x, r)``, where `x` is the root, and `r` is - a `RootResults` object. - disp : bool, optional - If True, raise RuntimeError if the algorithm didn't converge. - Otherwise, the convergence status is recorded in any `RootResults` - return object. - - Returns - ------- - root : float - Root of `f` between `a` and `b`. - r : `RootResults` (present if ``full_output = True``) - Object containing information about the convergence. In particular, - ``r.converged`` is True if the routine converged. - - Notes - ----- - `f` must be continuous. f(a) and f(b) must have opposite signs. - - Related functions fall into several classes: - - multivariate local optimizers - `fmin`, `fmin_powell`, `fmin_cg`, `fmin_bfgs`, `fmin_ncg` - nonlinear least squares minimizer - `leastsq` - constrained multivariate optimizers - `fmin_l_bfgs_b`, `fmin_tnc`, `fmin_cobyla` - global optimizers - `basinhopping`, `brute`, `differential_evolution` - local scalar minimizers - `fminbound`, `brent`, `golden`, `bracket` - N-D root-finding - `fsolve` - 1-D root-finding - `brenth`, `ridder`, `bisect`, `newton` - scalar fixed-point finder - `fixed_point` - - References - ---------- - .. [Brent1973] - Brent, R. P., - *Algorithms for Minimization Without Derivatives*. - Englewood Cliffs, NJ: Prentice-Hall, 1973. Ch. 3-4. - - .. [PressEtal1992] - Press, W. H.; Flannery, B. P.; Teukolsky, S. A.; and Vetterling, W. T. - *Numerical Recipes in FORTRAN: The Art of Scientific Computing*, 2nd ed. - Cambridge, England: Cambridge University Press, pp. 352-355, 1992. - Section 9.3: "Van Wijngaarden-Dekker-Brent Method." - - Examples - -------- - >>> def f(x): - ... return (x**2 - 1) - - >>> from scipy import optimize - - >>> root = optimize.brentq(f, -2, 0) - >>> root - -1.0 - - >>> root = optimize.brentq(f, 0, 2) - >>> root - 1.0 + Loosely adapted from + https://github.com/scipy/scipy/blob/d23363809572e9a44074a3f06f66137083446b48/scipy/optimize/_zeros_py.py#L682 """ maxiter = operator.index(maxiter) if xtol <= 0: raise ValueError("xtol too small (%g <= 0)" % xtol) - if rtol < _rtol: - raise ValueError(f"rtol too small ({rtol:g} < {_rtol:g})") + if rtol < BRENTQ_RTOL: + raise ValueError(f"rtol too small ({rtol:g} < {BRENTQ_RTOL:g})") f = _wrap_nan_raise(f) r = brentq_sf(f, a, b, xtol, rtol, maxiter) return r diff --git a/src/hapsira/core/math/ivp/_solver.py b/src/hapsira/core/math/ivp/_solver.py deleted file mode 100644 index ce68a76e7..000000000 --- a/src/hapsira/core/math/ivp/_solver.py +++ /dev/null @@ -1,131 +0,0 @@ -from numba import njit, jit - -from math import fabs - - -CONVERGED = 0 -SIGNERR = -1 -CONVERR = -2 -# EVALUEERR = -3 -INPROGRESS = 1 - - -@njit -def _min_hf(a, b): - return a if a < b else b - - -@njit -def _signbit_hf(a): - return a < 0 - - -@jit -def _brentq_hf( - func, # callback_type - xa, # double - xb, # double - xtol, # double - rtol, # double - iter_, # int -): - xpre, xcur = xa, xb - xblk = 0.0 - fpre, fcur, fblk = 0.0, 0.0, 0.0 - spre, scur = 0.0, 0.0 - - iterations = 0 - - fpre = func(xpre) - fcur = func(xcur) - funcalls = 2 - if fpre == 0: - return xpre, funcalls, iterations, CONVERGED - if fcur == 0: - return xcur, funcalls, iterations, CONVERGED - if _signbit_hf(fpre) == _signbit_hf(fcur): - return 0.0, funcalls, iterations, SIGNERR - - iterations = 0 - for _ in range(0, iter_): - iterations += 1 - if fpre != 0 and fcur != 0 and _signbit_hf(fpre) != _signbit_hf(fcur): - xblk = xpre - fblk = fpre - scur = xcur - xpre - spre = scur - if fabs(fblk) < fabs(fcur): - xpre = xcur - xcur = xblk - xblk = xpre - - fpre = fcur - fcur = fblk - fblk = fpre - - delta = (xtol + rtol * fabs(xcur)) / 2 - sbis = (xblk - xcur) / 2 - if fcur == 0 or fabs(sbis) < delta: - return xcur, funcalls, iterations, CONVERGED - - if fabs(spre) > delta and fabs(fcur) < fabs(fpre): - if xpre == xblk: - stry = -fcur * (xcur - xpre) / (fcur - fpre) - else: - dpre = (fpre - fcur) / (xpre - xcur) - dblk = (fblk - fcur) / (xblk - xcur) - stry = ( - -fcur * (fblk * dblk - fpre * dpre) / (dblk * dpre * (fblk - fpre)) - ) - if 2 * fabs(stry) < _min_hf(fabs(spre), 3 * fabs(sbis) - delta): - spre = scur - scur = stry - else: - spre = sbis - scur = sbis - else: - spre = sbis - scur = sbis - - xpre = xcur - fpre = fcur - if fabs(scur) > delta: - xcur += scur - else: - xcur += delta if sbis > 0 else -delta - - fcur = func(xcur) - funcalls += 1 - - return xcur, funcalls, iterations, CONVERR - - -@jit -def brentq_sf( - func, # func - a, # double - b, # double - xtol, # double - rtol, # double - iter_, # int -): - if xtol < 0: - raise ValueError("xtol must be >= 0") - if iter_ < 0: - raise ValueError("maxiter should be > 0") - - zero, funcalls, iterations, error_num = _brentq_hf( - func, - a, - b, - xtol, - rtol, - iter_, - ) - - if error_num == SIGNERR: - raise ValueError("f(a) and f(b) must have different signs") - if error_num == CONVERR: - raise RuntimeError("Failed to converge after %d iterations." % iterations) - - return zero # double From c3659f5501589e4f64d01eea9e55dea3a59b91b5 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 13 Jan 2024 08:15:51 +0100 Subject: [PATCH 092/346] cleanup --- src/hapsira/core/math/ivp/_api.py | 55 ++----------------------------- 1 file changed, 2 insertions(+), 53 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 9917e24af..2517e6f07 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -17,8 +17,6 @@ class OdeResult(dict): Attributes ---------- - x : ndarray - The solution of the optimization. success : bool Whether or not the optimizer exited successfully. status : int @@ -26,29 +24,10 @@ class OdeResult(dict): underlying solver. Refer to `message` for details. message : str Description of the cause of the termination. - fun, jac, hess: ndarray - Values of objective function, its Jacobian and its Hessian (if - available). The Hessians may be approximations, see the documentation - of the function in question. - hess_inv : object - Inverse of the objective function's Hessian; may be an approximation. - Not available for all solvers. The type of this attribute may be - either np.ndarray or scipy.sparse.linalg.LinearOperator. - nfev, njev, nhev : int + nfev, njev : int Number of evaluations of the objective functions and of its Jacobian and Hessian. - nit : int - Number of iterations performed by the optimizer. - maxcv : float - The maximum constraint violation. - - Notes - ----- - Depending on the specific solver being used, `OptimizeResult` may - not have all attributes listed here, and they may have additional - attributes not listed here. Since this class is essentially a - subclass of dict with attribute accessors, one can see which - attributes are available using the `OptimizeResult.keys` method. + t, y, sol, t_events, y_events, nlu : ? """ def __getattr__(self, name): @@ -60,36 +39,6 @@ def __getattr__(self, name): __setattr__ = dict.__setitem__ __delattr__ = dict.__delitem__ - # def __repr__(self): - # order_keys = ['message', 'success', 'status', 'fun', 'funl', 'x', 'xl', - # 'col_ind', 'nit', 'lower', 'upper', 'eqlin', 'ineqlin', - # 'converged', 'flag', 'function_calls', 'iterations', - # 'root'] - # order_keys = getattr(self, '_order_keys', order_keys) - # # 'slack', 'con' are redundant with residuals - # # 'crossover_nit' is probably not interesting to most users - # omit_keys = {'slack', 'con', 'crossover_nit', '_order_keys'} - - # def key(item): - # try: - # return order_keys.index(item[0].lower()) - # except ValueError: # item not in list - # return np.inf - - # def omit_redundant(items): - # for item in items: - # if item[0] in omit_keys: - # continue - # yield item - - # def item_sorter(d): - # return sorted(omit_redundant(d.items()), key=key) - - # if self.keys(): - # return _dict_formatter(self, sorter=item_sorter) - # else: - # return self.__class__.__name__ + "()" - def __dir__(self): return list(self.keys()) From eb76742fe65db60216057066b790d718a12a3442 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 13 Jan 2024 08:48:00 +0100 Subject: [PATCH 093/346] rm dead code --- src/hapsira/core/math/ivp/_api.py | 122 ++++-------------------------- 1 file changed, 15 insertions(+), 107 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 2517e6f07..4f0afd6e9 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -175,8 +175,8 @@ def solve_ivp( t_span, y0, method=DOP853, - t_eval=None, - dense_output=False, + # t_eval=None, + # dense_output=False, events=None, vectorized=False, args=None, @@ -208,11 +208,6 @@ def solve_ivp( Can be applied in the complex domain. You can also pass an arbitrary class derived from `OdeSolver` which implements the solver. - t_eval : array_like or None, optional - Times at which to store the computed solution, must be sorted and lie - within `t_span`. If None (default), use points selected by the solver. - dense_output : bool, optional - Whether to compute a continuous solution. Default is False. events : callable, or list of callables, optional Events to track. If None (default), no events will be tracked. Each event occurs at the zeros of a continuous function of time and @@ -356,68 +351,18 @@ def solve_ivp( occurred (``status >= 0``). """ - # if method not in METHODS and not ( - # inspect.isclass(method) and issubclass(method, OdeSolver)): - # raise ValueError("`method` must be one of {} or OdeSolver class." - # .format(METHODS)) t0, tf = map(float, t_span) - if args is not None: - # Wrap the user's fun (and jac, if given) in lambdas to hide the - # additional parameters. Pass in the original fun as a keyword - # argument to keep it in the scope of the lambda. - try: - _ = [*(args)] - except TypeError as exp: - suggestion_tuple = ( - "Supplied 'args' cannot be unpacked. Please supply `args`" - f" as a tuple (e.g. `args=({args},)`)" - ) - raise TypeError(suggestion_tuple) from exp - - def fun(t, x, fun=fun): - return fun(t, x, *args) - - jac = options.get("jac") - if callable(jac): - options["jac"] = lambda t, x: jac(t, x, *args) - - if t_eval is not None: - t_eval = np.asarray(t_eval) - if t_eval.ndim != 1: - raise ValueError("`t_eval` must be 1-dimensional.") - - if np.any(t_eval < min(t0, tf)) or np.any(t_eval > max(t0, tf)): - raise ValueError("Values in `t_eval` are not within `t_span`.") - - d = np.diff(t_eval) - if tf > t0 and np.any(d <= 0) or tf < t0 and np.any(d >= 0): - raise ValueError("Values in `t_eval` are not properly sorted.") - - if tf > t0: - t_eval_i = 0 - else: - # Make order of t_eval decreasing to use np.searchsorted. - t_eval = t_eval[::-1] - # This will be an upper bound for slices. - t_eval_i = t_eval.shape[0] + assert isinstance(args, tuple) - # if method in METHODS: - # method = METHODS[method] + def fun(t, x, fun=fun): + return fun(t, x, *args) solver = method(fun, t0, y0, tf, vectorized=vectorized, **options) - if t_eval is None: - ts = [t0] - ys = [y0] - elif t_eval is not None and dense_output: - ts = [] - ti = [t0] - ys = [] - else: - ts = [] - ys = [] + ts = [t0] + ys = [y0] interpolants = [] @@ -451,11 +396,8 @@ def fun(t, x, fun=fun): t = solver.t y = solver.y - if dense_output: - sol = solver.dense_output() - interpolants.append(sol) - else: - sol = None + sol = solver.dense_output() + interpolants.append(sol) if events is not None: g_new = [event(t, y) for event in events] @@ -479,51 +421,17 @@ def fun(t, x, fun=fun): g = g_new - if t_eval is None: - ts.append(t) - ys.append(y) - else: - # The value in t_eval equal to t will be included. - if solver.direction > 0: - t_eval_i_new = np.searchsorted(t_eval, t, side="right") - t_eval_step = t_eval[t_eval_i:t_eval_i_new] - else: - t_eval_i_new = np.searchsorted(t_eval, t, side="left") - # It has to be done with two slice operations, because - # you can't slice to 0th element inclusive using backward - # slicing. - t_eval_step = t_eval[t_eval_i_new:t_eval_i][::-1] - - if t_eval_step.size > 0: - if sol is None: - sol = solver.dense_output() - ts.append(t_eval_step) - ys.append(sol(t_eval_step)) - t_eval_i = t_eval_i_new - - if t_eval is not None and dense_output: - ti.append(t) - - message = MESSAGES.get(status, message) + ts.append(t) + ys.append(y) if t_events is not None: t_events = [np.asarray(te) for te in t_events] y_events = [np.asarray(ye) for ye in y_events] - if t_eval is None: - ts = np.array(ts) - ys = np.vstack(ys).T - elif ts: - ts = np.hstack(ts) - ys = np.hstack(ys) + ts = np.array(ts) + ys = np.vstack(ys).T - if dense_output: - if t_eval is None: - sol = OdeSolution(ts, interpolants) - else: - sol = OdeSolution(ti, interpolants) - else: - sol = None + sol = OdeSolution(ts, interpolants) return OdeResult( t=ts, @@ -535,6 +443,6 @@ def fun(t, x, fun=fun): njev=solver.njev, nlu=solver.nlu, status=status, - message=message, + message=MESSAGES.get(status, message), success=status >= 0, ) From 8f3e2415e1713fd80b822fc3ea470c78e8c9b69a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 13 Jan 2024 09:01:26 +0100 Subject: [PATCH 094/346] rm dead code --- src/hapsira/core/math/ivp/_common.py | 64 ++++------------------------ 1 file changed, 9 insertions(+), 55 deletions(-) diff --git a/src/hapsira/core/math/ivp/_common.py b/src/hapsira/core/math/ivp/_common.py index 692e7eb71..a78205d62 100644 --- a/src/hapsira/core/math/ivp/_common.py +++ b/src/hapsira/core/math/ivp/_common.py @@ -1,5 +1,3 @@ -from itertools import groupby - import numpy as np @@ -46,30 +44,12 @@ def __init__(self, ts, interpolants): self.ts = ts self.interpolants = interpolants - if ts[-1] >= ts[0]: - self.t_min = ts[0] - self.t_max = ts[-1] - self.ascending = True - self.ts_sorted = ts - else: - self.t_min = ts[-1] - self.t_max = ts[0] - self.ascending = False - self.ts_sorted = ts[::-1] - - def _call_single(self, t): - # Here we preserve a certain symmetry that when t is in self.ts, - # then we prioritize a segment with a lower index. - if self.ascending: - ind = np.searchsorted(self.ts_sorted, t, side="left") - else: - ind = np.searchsorted(self.ts_sorted, t, side="right") - segment = min(max(ind - 1, 0), self.n_segments - 1) - if not self.ascending: - segment = self.n_segments - 1 - segment + assert ts[-1] >= ts[0] - return self.interpolants[segment](t) + self.t_min = ts[0] + self.t_max = ts[-1] + self.ts_sorted = ts def __call__(self, t): """Evaluate the solution. @@ -87,34 +67,8 @@ def __call__(self, t): """ t = np.asarray(t) - if t.ndim == 0: - return self._call_single(t) - - order = np.argsort(t) - reverse = np.empty_like(order) - reverse[order] = np.arange(order.shape[0]) - t_sorted = t[order] - - # See comment in self._call_single. - if self.ascending: - segments = np.searchsorted(self.ts_sorted, t_sorted, side="left") - else: - segments = np.searchsorted(self.ts_sorted, t_sorted, side="right") - segments -= 1 - segments[segments < 0] = 0 - segments[segments > self.n_segments - 1] = self.n_segments - 1 - if not self.ascending: - segments = self.n_segments - 1 - segments - - ys = [] - group_start = 0 - for segment, group in groupby(segments): - group_end = group_start + len(list(group)) - y = self.interpolants[segment](t_sorted[group_start:group_end]) - ys.append(y) - group_start = group_end - - ys = np.hstack(ys) - ys = ys[:, reverse] - - return ys + assert t.ndim == 0 + + ind = np.searchsorted(self.ts_sorted, t, side="left") + segment = min(max(ind - 1, 0), self.n_segments - 1) + return self.interpolants[segment](t) From 243979e1b1c44b54555b2fcd8ea8a7b8f0b82178 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 13 Jan 2024 09:12:14 +0100 Subject: [PATCH 095/346] rm dead code --- src/hapsira/core/math/ivp/_base.py | 67 +++++++----------------------- 1 file changed, 15 insertions(+), 52 deletions(-) diff --git a/src/hapsira/core/math/ivp/_base.py b/src/hapsira/core/math/ivp/_base.py index a19084e33..f984cfa2a 100644 --- a/src/hapsira/core/math/ivp/_base.py +++ b/src/hapsira/core/math/ivp/_base.py @@ -1,25 +1,18 @@ import numpy as np -def check_arguments(fun, y0, support_complex): +def check_arguments(fun, y0): """Helper function for checking arguments common to all solvers.""" + y0 = np.asarray(y0) - if np.issubdtype(y0.dtype, np.complexfloating): - if not support_complex: - raise ValueError( - "`y0` is complex, but the chosen solver does " - "not support integration in a complex domain." - ) - dtype = complex - else: - dtype = float - y0 = y0.astype(dtype, copy=False) - if y0.ndim != 1: - raise ValueError("`y0` must be 1-dimensional.") + assert not np.issubdtype(y0.dtype, np.complexfloating) + + dtype = float + y0 = y0.astype(dtype, copy=False) - if not np.isfinite(y0).all(): - raise ValueError("All components of the initial state `y0` must be finite.") + assert not y0.ndim != 1 + assert np.isfinite(y0).all() def fun_wrapped(t, y): return np.asarray(fun(t, y), dtype=dtype) @@ -98,10 +91,6 @@ class OdeSolver: will result in slower execution for other methods. It can also result in slower overall execution for 'Radau' and 'BDF' in some circumstances (e.g. small ``len(y0)``). - support_complex : bool, optional - Whether integration in a complex domain should be supported. - Generally determined by a derived solver class capabilities. - Default is False. Attributes ---------- @@ -131,10 +120,10 @@ class OdeSolver: TOO_SMALL_STEP = "Required step size is less than spacing between numbers." - def __init__(self, fun, t0, y0, t_bound, vectorized, support_complex=False): + def __init__(self, fun, t0, y0, t_bound, vectorized): self.t_old = None self.t = t0 - self._fun, self.y = check_arguments(fun, y0, support_complex) + self._fun, self.y = check_arguments(fun, y0) self.t_bound = t_bound self.vectorized = vectorized @@ -216,16 +205,11 @@ def dense_output(self): sol : `DenseOutput` Local interpolant over the last successful step. """ - if self.t_old is None: - raise RuntimeError( - "Dense output is available after a successful " "step was made." - ) + assert self.t_old is not None - if self.n == 0 or self.t == self.t_old: - # Handle corner cases of empty solver and no integration. - return ConstantDenseOutput(self.t_old, self.t, self.y) - else: - return self._dense_output_impl() + assert not (self.n == 0 or self.t == self.t_old) + + return self._dense_output_impl() def _step_impl(self): raise NotImplementedError @@ -268,29 +252,8 @@ def __call__(self, t): 1-D array. """ t = np.asarray(t) - if t.ndim > 1: - raise ValueError("`t` must be a float or a 1-D array.") + assert not t.ndim > 1 return self._call_impl(t) def _call_impl(self, t): raise NotImplementedError - - -class ConstantDenseOutput(DenseOutput): - """Constant value interpolator. - - This class used for degenerate integration cases: equal integration limits - or a system with 0 equations. - """ - - def __init__(self, t_old, t, value): - super().__init__(t_old, t) - self.value = value - - def _call_impl(self, t): - if t.ndim == 0: - return self.value - else: - ret = np.empty((self.value.shape[0], t.shape[0])) - ret[:] = self.value[:, None] - return ret From 2509e9fcc9f6c932220b413019bad9044368539d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 13 Jan 2024 09:18:18 +0100 Subject: [PATCH 096/346] mv code --- src/hapsira/core/math/ivp/_base.py | 259 ---------------------------- src/hapsira/core/math/ivp/_rk.py | 261 ++++++++++++++++++++++++++++- 2 files changed, 259 insertions(+), 261 deletions(-) diff --git a/src/hapsira/core/math/ivp/_base.py b/src/hapsira/core/math/ivp/_base.py index f984cfa2a..e69de29bb 100644 --- a/src/hapsira/core/math/ivp/_base.py +++ b/src/hapsira/core/math/ivp/_base.py @@ -1,259 +0,0 @@ -import numpy as np - - -def check_arguments(fun, y0): - """Helper function for checking arguments common to all solvers.""" - - y0 = np.asarray(y0) - - assert not np.issubdtype(y0.dtype, np.complexfloating) - - dtype = float - y0 = y0.astype(dtype, copy=False) - - assert not y0.ndim != 1 - assert np.isfinite(y0).all() - - def fun_wrapped(t, y): - return np.asarray(fun(t, y), dtype=dtype) - - return fun_wrapped, y0 - - -class OdeSolver: - """Base class for ODE solvers. - - In order to implement a new solver you need to follow the guidelines: - - 1. A constructor must accept parameters presented in the base class - (listed below) along with any other parameters specific to a solver. - 2. A constructor must accept arbitrary extraneous arguments - ``**extraneous``, but warn that these arguments are irrelevant - using `common.warn_extraneous` function. Do not pass these - arguments to the base class. - 3. A solver must implement a private method `_step_impl(self)` which - propagates a solver one step further. It must return tuple - ``(success, message)``, where ``success`` is a boolean indicating - whether a step was successful, and ``message`` is a string - containing description of a failure if a step failed or None - otherwise. - 4. A solver must implement a private method `_dense_output_impl(self)`, - which returns a `DenseOutput` object covering the last successful - step. - 5. A solver must have attributes listed below in Attributes section. - Note that ``t_old`` and ``step_size`` are updated automatically. - 6. Use `fun(self, t, y)` method for the system rhs evaluation, this - way the number of function evaluations (`nfev`) will be tracked - automatically. - 7. For convenience, a base class provides `fun_single(self, t, y)` and - `fun_vectorized(self, t, y)` for evaluating the rhs in - non-vectorized and vectorized fashions respectively (regardless of - how `fun` from the constructor is implemented). These calls don't - increment `nfev`. - 8. If a solver uses a Jacobian matrix and LU decompositions, it should - track the number of Jacobian evaluations (`njev`) and the number of - LU decompositions (`nlu`). - 9. By convention, the function evaluations used to compute a finite - difference approximation of the Jacobian should not be counted in - `nfev`, thus use `fun_single(self, t, y)` or - `fun_vectorized(self, t, y)` when computing a finite difference - approximation of the Jacobian. - - Parameters - ---------- - fun : callable - Right-hand side of the system: the time derivative of the state ``y`` - at time ``t``. The calling signature is ``fun(t, y)``, where ``t`` is a - scalar and ``y`` is an ndarray with ``len(y) = len(y0)``. ``fun`` must - return an array of the same shape as ``y``. See `vectorized` for more - information. - t0 : float - Initial time. - y0 : array_like, shape (n,) - Initial state. - t_bound : float - Boundary time --- the integration won't continue beyond it. It also - determines the direction of the integration. - vectorized : bool - Whether `fun` can be called in a vectorized fashion. Default is False. - - If ``vectorized`` is False, `fun` will always be called with ``y`` of - shape ``(n,)``, where ``n = len(y0)``. - - If ``vectorized`` is True, `fun` may be called with ``y`` of shape - ``(n, k)``, where ``k`` is an integer. In this case, `fun` must behave - such that ``fun(t, y)[:, i] == fun(t, y[:, i])`` (i.e. each column of - the returned array is the time derivative of the state corresponding - with a column of ``y``). - - Setting ``vectorized=True`` allows for faster finite difference - approximation of the Jacobian by methods 'Radau' and 'BDF', but - will result in slower execution for other methods. It can also - result in slower overall execution for 'Radau' and 'BDF' in some - circumstances (e.g. small ``len(y0)``). - - Attributes - ---------- - n : int - Number of equations. - status : string - Current status of the solver: 'running', 'finished' or 'failed'. - t_bound : float - Boundary time. - direction : float - Integration direction: +1 or -1. - t : float - Current time. - y : ndarray - Current state. - t_old : float - Previous time. None if no steps were made yet. - step_size : float - Size of the last successful step. None if no steps were made yet. - nfev : int - Number of the system's rhs evaluations. - njev : int - Number of the Jacobian evaluations. - nlu : int - Number of LU decompositions. - """ - - TOO_SMALL_STEP = "Required step size is less than spacing between numbers." - - def __init__(self, fun, t0, y0, t_bound, vectorized): - self.t_old = None - self.t = t0 - self._fun, self.y = check_arguments(fun, y0) - self.t_bound = t_bound - self.vectorized = vectorized - - if vectorized: - - def fun_single(t, y): - return self._fun(t, y[:, None]).ravel() - - fun_vectorized = self._fun - else: - fun_single = self._fun - - def fun_vectorized(t, y): - f = np.empty_like(y) - for i, yi in enumerate(y.T): - f[:, i] = self._fun(t, yi) - return f - - def fun(t, y): - self.nfev += 1 - return self.fun_single(t, y) - - self.fun = fun - self.fun_single = fun_single - self.fun_vectorized = fun_vectorized - - self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 - self.n = self.y.size - self.status = "running" - - self.nfev = 0 - self.njev = 0 - self.nlu = 0 - - @property - def step_size(self): - if self.t_old is None: - return None - else: - return np.abs(self.t - self.t_old) - - def step(self): - """Perform one integration step. - - Returns - ------- - message : string or None - Report from the solver. Typically a reason for a failure if - `self.status` is 'failed' after the step was taken or None - otherwise. - """ - if self.status != "running": - raise RuntimeError("Attempt to step on a failed or finished " "solver.") - - if self.n == 0 or self.t == self.t_bound: - # Handle corner cases of empty solver or no integration. - self.t_old = self.t - self.t = self.t_bound - message = None - self.status = "finished" - else: - t = self.t - success, message = self._step_impl() - - if not success: - self.status = "failed" - else: - self.t_old = t - if self.direction * (self.t - self.t_bound) >= 0: - self.status = "finished" - - return message - - def dense_output(self): - """Compute a local interpolant over the last successful step. - - Returns - ------- - sol : `DenseOutput` - Local interpolant over the last successful step. - """ - assert self.t_old is not None - - assert not (self.n == 0 or self.t == self.t_old) - - return self._dense_output_impl() - - def _step_impl(self): - raise NotImplementedError - - def _dense_output_impl(self): - raise NotImplementedError - - -class DenseOutput: - """Base class for local interpolant over step made by an ODE solver. - - It interpolates between `t_min` and `t_max` (see Attributes below). - Evaluation outside this interval is not forbidden, but the accuracy is not - guaranteed. - - Attributes - ---------- - t_min, t_max : float - Time range of the interpolation. - """ - - def __init__(self, t_old, t): - self.t_old = t_old - self.t = t - self.t_min = min(t, t_old) - self.t_max = max(t, t_old) - - def __call__(self, t): - """Evaluate the interpolant. - - Parameters - ---------- - t : float or array_like with shape (n_points,) - Points to evaluate the solution at. - - Returns - ------- - y : ndarray, shape (n,) or (n, n_points) - Computed values. Shape depends on whether `t` was a scalar or a - 1-D array. - """ - t = np.asarray(t) - assert not t.ndim > 1 - return self._call_impl(t) - - def _call_impl(self, t): - raise NotImplementedError diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 206558596..828917e50 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -3,7 +3,6 @@ import numpy as np from . import _dop853_coefficients as dop853_coefficients -from ._base import DenseOutput, OdeSolver EPS = np.finfo(float).eps @@ -194,6 +193,264 @@ def validate_tol(rtol, atol, n): return rtol, atol +class DenseOutput: + """Base class for local interpolant over step made by an ODE solver. + + It interpolates between `t_min` and `t_max` (see Attributes below). + Evaluation outside this interval is not forbidden, but the accuracy is not + guaranteed. + + Attributes + ---------- + t_min, t_max : float + Time range of the interpolation. + """ + + def __init__(self, t_old, t): + self.t_old = t_old + self.t = t + self.t_min = min(t, t_old) + self.t_max = max(t, t_old) + + def __call__(self, t): + """Evaluate the interpolant. + + Parameters + ---------- + t : float or array_like with shape (n_points,) + Points to evaluate the solution at. + + Returns + ------- + y : ndarray, shape (n,) or (n, n_points) + Computed values. Shape depends on whether `t` was a scalar or a + 1-D array. + """ + t = np.asarray(t) + assert not t.ndim > 1 + return self._call_impl(t) + + def _call_impl(self, t): + raise NotImplementedError + + +def check_arguments(fun, y0): + """Helper function for checking arguments common to all solvers.""" + + y0 = np.asarray(y0) + + assert not np.issubdtype(y0.dtype, np.complexfloating) + + dtype = float + y0 = y0.astype(dtype, copy=False) + + assert not y0.ndim != 1 + assert np.isfinite(y0).all() + + def fun_wrapped(t, y): + return np.asarray(fun(t, y), dtype=dtype) + + return fun_wrapped, y0 + + +class OdeSolver: + """Base class for ODE solvers. + + In order to implement a new solver you need to follow the guidelines: + + 1. A constructor must accept parameters presented in the base class + (listed below) along with any other parameters specific to a solver. + 2. A constructor must accept arbitrary extraneous arguments + ``**extraneous``, but warn that these arguments are irrelevant + using `common.warn_extraneous` function. Do not pass these + arguments to the base class. + 3. A solver must implement a private method `_step_impl(self)` which + propagates a solver one step further. It must return tuple + ``(success, message)``, where ``success`` is a boolean indicating + whether a step was successful, and ``message`` is a string + containing description of a failure if a step failed or None + otherwise. + 4. A solver must implement a private method `_dense_output_impl(self)`, + which returns a `DenseOutput` object covering the last successful + step. + 5. A solver must have attributes listed below in Attributes section. + Note that ``t_old`` and ``step_size`` are updated automatically. + 6. Use `fun(self, t, y)` method for the system rhs evaluation, this + way the number of function evaluations (`nfev`) will be tracked + automatically. + 7. For convenience, a base class provides `fun_single(self, t, y)` and + `fun_vectorized(self, t, y)` for evaluating the rhs in + non-vectorized and vectorized fashions respectively (regardless of + how `fun` from the constructor is implemented). These calls don't + increment `nfev`. + 8. If a solver uses a Jacobian matrix and LU decompositions, it should + track the number of Jacobian evaluations (`njev`) and the number of + LU decompositions (`nlu`). + 9. By convention, the function evaluations used to compute a finite + difference approximation of the Jacobian should not be counted in + `nfev`, thus use `fun_single(self, t, y)` or + `fun_vectorized(self, t, y)` when computing a finite difference + approximation of the Jacobian. + + Parameters + ---------- + fun : callable + Right-hand side of the system: the time derivative of the state ``y`` + at time ``t``. The calling signature is ``fun(t, y)``, where ``t`` is a + scalar and ``y`` is an ndarray with ``len(y) = len(y0)``. ``fun`` must + return an array of the same shape as ``y``. See `vectorized` for more + information. + t0 : float + Initial time. + y0 : array_like, shape (n,) + Initial state. + t_bound : float + Boundary time --- the integration won't continue beyond it. It also + determines the direction of the integration. + vectorized : bool + Whether `fun` can be called in a vectorized fashion. Default is False. + + If ``vectorized`` is False, `fun` will always be called with ``y`` of + shape ``(n,)``, where ``n = len(y0)``. + + If ``vectorized`` is True, `fun` may be called with ``y`` of shape + ``(n, k)``, where ``k`` is an integer. In this case, `fun` must behave + such that ``fun(t, y)[:, i] == fun(t, y[:, i])`` (i.e. each column of + the returned array is the time derivative of the state corresponding + with a column of ``y``). + + Setting ``vectorized=True`` allows for faster finite difference + approximation of the Jacobian by methods 'Radau' and 'BDF', but + will result in slower execution for other methods. It can also + result in slower overall execution for 'Radau' and 'BDF' in some + circumstances (e.g. small ``len(y0)``). + + Attributes + ---------- + n : int + Number of equations. + status : string + Current status of the solver: 'running', 'finished' or 'failed'. + t_bound : float + Boundary time. + direction : float + Integration direction: +1 or -1. + t : float + Current time. + y : ndarray + Current state. + t_old : float + Previous time. None if no steps were made yet. + step_size : float + Size of the last successful step. None if no steps were made yet. + nfev : int + Number of the system's rhs evaluations. + njev : int + Number of the Jacobian evaluations. + nlu : int + Number of LU decompositions. + """ + + TOO_SMALL_STEP = "Required step size is less than spacing between numbers." + + def __init__(self, fun, t0, y0, t_bound, vectorized): + self.t_old = None + self.t = t0 + self._fun, self.y = check_arguments(fun, y0) + self.t_bound = t_bound + self.vectorized = vectorized + + if vectorized: + + def fun_single(t, y): + return self._fun(t, y[:, None]).ravel() + + fun_vectorized = self._fun + else: + fun_single = self._fun + + def fun_vectorized(t, y): + f = np.empty_like(y) + for i, yi in enumerate(y.T): + f[:, i] = self._fun(t, yi) + return f + + def fun(t, y): + self.nfev += 1 + return self.fun_single(t, y) + + self.fun = fun + self.fun_single = fun_single + self.fun_vectorized = fun_vectorized + + self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 + self.n = self.y.size + self.status = "running" + + self.nfev = 0 + self.njev = 0 + self.nlu = 0 + + @property + def step_size(self): + if self.t_old is None: + return None + else: + return np.abs(self.t - self.t_old) + + def step(self): + """Perform one integration step. + + Returns + ------- + message : string or None + Report from the solver. Typically a reason for a failure if + `self.status` is 'failed' after the step was taken or None + otherwise. + """ + if self.status != "running": + raise RuntimeError("Attempt to step on a failed or finished " "solver.") + + if self.n == 0 or self.t == self.t_bound: + # Handle corner cases of empty solver or no integration. + self.t_old = self.t + self.t = self.t_bound + message = None + self.status = "finished" + else: + t = self.t + success, message = self._step_impl() + + if not success: + self.status = "failed" + else: + self.t_old = t + if self.direction * (self.t - self.t_bound) >= 0: + self.status = "finished" + + return message + + def dense_output(self): + """Compute a local interpolant over the last successful step. + + Returns + ------- + sol : `DenseOutput` + Local interpolant over the last successful step. + """ + assert self.t_old is not None + + assert not (self.n == 0 or self.t == self.t_old) + + return self._dense_output_impl() + + def _step_impl(self): + raise NotImplementedError + + def _dense_output_impl(self): + raise NotImplementedError + + class RungeKutta(OdeSolver): """Base class for explicit Runge-Kutta methods.""" @@ -220,7 +477,7 @@ def __init__( **extraneous, ): warn_extraneous(extraneous) - super().__init__(fun, t0, y0, t_bound, vectorized, support_complex=True) + super().__init__(fun, t0, y0, t_bound, vectorized) self.y_old = None self.max_step = validate_max_step(max_step) self.rtol, self.atol = validate_tol(rtol, atol, self.n) From 3d7b66bdc737d22d7ded7f9112f2a63322be79a9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 13 Jan 2024 09:18:42 +0100 Subject: [PATCH 097/346] rm empty file --- src/hapsira/core/math/ivp/_base.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 src/hapsira/core/math/ivp/_base.py diff --git a/src/hapsira/core/math/ivp/_base.py b/src/hapsira/core/math/ivp/_base.py deleted file mode 100644 index e69de29bb..000000000 From fab3b44b308e190bfc5572957f63dd774822de8b Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 13 Jan 2024 09:28:12 +0100 Subject: [PATCH 098/346] rm dead code --- src/hapsira/core/math/ivp/_rk.py | 87 +++++++++----------------------- 1 file changed, 23 insertions(+), 64 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 828917e50..164bb01f4 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -193,12 +193,8 @@ def validate_tol(rtol, atol, n): return rtol, atol -class DenseOutput: - """Base class for local interpolant over step made by an ODE solver. - - It interpolates between `t_min` and `t_max` (see Attributes below). - Evaluation outside this interval is not forbidden, but the accuracy is not - guaranteed. +class Dop853DenseOutput: + """local interpolant over step made by an ODE solver. Attributes ---------- @@ -206,11 +202,14 @@ class DenseOutput: Time range of the interpolation. """ - def __init__(self, t_old, t): + def __init__(self, t_old, t, y_old, F): self.t_old = t_old self.t = t self.t_min = min(t, t_old) self.t_max = max(t, t_old) + self.h = t - t_old + self.F = F + self.y_old = y_old def __call__(self, t): """Evaluate the interpolant. @@ -231,7 +230,23 @@ def __call__(self, t): return self._call_impl(t) def _call_impl(self, t): - raise NotImplementedError + x = (t - self.t_old) / self.h + + if t.ndim == 0: + y = np.zeros_like(self.y_old) + else: + x = x[:, None] + y = np.zeros((len(x), len(self.y_old)), dtype=self.y_old.dtype) + + for i, f in enumerate(reversed(self.F)): + y += f + if i % 2 == 0: + y *= x + else: + y *= 1 - x + y += self.y_old + + return y.T def check_arguments(fun, y0): @@ -571,10 +586,6 @@ def _step_impl(self): return True, None - def _dense_output_impl(self): - Q = self.K.T.dot(self.P) - return RkDenseOutput(self.t_old, self.t, self.y_old, Q) - class DOP853(RungeKutta): """Explicit Runge-Kutta method of order 8. @@ -743,55 +754,3 @@ def _dense_output_impl(self): F[3:] = h * np.dot(self.D, K) return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) - - -class Dop853DenseOutput(DenseOutput): - def __init__(self, t_old, t, y_old, F): - super().__init__(t_old, t) - self.h = t - t_old - self.F = F - self.y_old = y_old - - def _call_impl(self, t): - x = (t - self.t_old) / self.h - - if t.ndim == 0: - y = np.zeros_like(self.y_old) - else: - x = x[:, None] - y = np.zeros((len(x), len(self.y_old)), dtype=self.y_old.dtype) - - for i, f in enumerate(reversed(self.F)): - y += f - if i % 2 == 0: - y *= x - else: - y *= 1 - x - y += self.y_old - - return y.T - - -class RkDenseOutput(DenseOutput): - def __init__(self, t_old, t, y_old, Q): - super().__init__(t_old, t) - self.h = t - t_old - self.Q = Q - self.order = Q.shape[1] - 1 - self.y_old = y_old - - def _call_impl(self, t): - x = (t - self.t_old) / self.h - if t.ndim == 0: - p = np.tile(x, self.order + 1) - p = np.cumprod(p) - else: - p = np.tile(x, (self.order + 1, 1)) - p = np.cumprod(p, axis=0) - y = self.h * np.dot(self.Q, p) - if y.ndim == 2: - y += self.y_old[:, None] - else: - y += self.y_old - - return y From 08cf5f022c614c6f10ac8aa046bd8ff4ca7f16e3 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 13 Jan 2024 09:44:20 +0100 Subject: [PATCH 099/346] rm class --- src/hapsira/core/math/ivp/_rk.py | 264 +++++++++++++------------------ 1 file changed, 108 insertions(+), 156 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 164bb01f4..3c12cc829 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -466,17 +466,23 @@ def _dense_output_impl(self): raise NotImplementedError -class RungeKutta(OdeSolver): +class DOP853(OdeSolver): """Base class for explicit Runge-Kutta methods.""" - C: np.ndarray = NotImplemented - A: np.ndarray = NotImplemented - B: np.ndarray = NotImplemented + n_stages: int = dop853_coefficients.N_STAGES + order: int = 8 + error_estimator_order: int = 7 + A = dop853_coefficients.A[:n_stages, :n_stages] + B = dop853_coefficients.B + C = dop853_coefficients.C[:n_stages] E: np.ndarray = NotImplemented + E3 = dop853_coefficients.E3 + E5 = dop853_coefficients.E5 + D = dop853_coefficients.D P: np.ndarray = NotImplemented - order: int = NotImplemented - error_estimator_order: int = NotImplemented - n_stages: int = NotImplemented + + A_EXTRA = dop853_coefficients.A[n_stages + 1 :] + C_EXTRA = dop853_coefficients.C[n_stages + 1 :] def __init__( self, @@ -514,11 +520,20 @@ def __init__( self.error_exponent = -1 / (self.error_estimator_order + 1) self.h_previous = None - def _estimate_error(self, K, h): - return np.dot(K.T, self.E) * h + self.K_extended = np.empty( + (dop853_coefficients.N_STAGES_EXTENDED, self.n), dtype=self.y.dtype + ) + self.K = self.K_extended[: self.n_stages + 1] def _estimate_error_norm(self, K, h, scale): - return norm(self._estimate_error(K, h) / scale) + err5 = np.dot(K.T, self.E5) / scale + err3 = np.dot(K.T, self.E3) / scale + err5_norm_2 = np.linalg.norm(err5) ** 2 + err3_norm_2 = np.linalg.norm(err3) ** 2 + if err5_norm_2 == 0 and err3_norm_2 == 0: + return 0.0 + denom = err5_norm_2 + 0.01 * err3_norm_2 + return np.abs(h) * err5_norm_2 / np.sqrt(denom * len(scale)) def _step_impl(self): t = self.t @@ -586,152 +601,6 @@ def _step_impl(self): return True, None - -class DOP853(RungeKutta): - """Explicit Runge-Kutta method of order 8. - - This is a Python implementation of "DOP853" algorithm originally written - in Fortran [1]_, [2]_. Note that this is not a literate translation, but - the algorithmic core and coefficients are the same. - - Can be applied in the complex domain. - - Parameters - ---------- - fun : callable - Right-hand side of the system. The calling signature is ``fun(t, y)``. - Here, ``t`` is a scalar, and there are two options for the ndarray ``y``: - It can either have shape (n,); then ``fun`` must return array_like with - shape (n,). Alternatively it can have shape (n, k); then ``fun`` - must return an array_like with shape (n, k), i.e. each column - corresponds to a single column in ``y``. The choice between the two - options is determined by `vectorized` argument (see below). - t0 : float - Initial time. - y0 : array_like, shape (n,) - Initial state. - t_bound : float - Boundary time - the integration won't continue beyond it. It also - determines the direction of the integration. - first_step : float or None, optional - Initial step size. Default is ``None`` which means that the algorithm - should choose. - max_step : float, optional - Maximum allowed step size. Default is np.inf, i.e. the step size is not - bounded and determined solely by the solver. - rtol, atol : float and array_like, optional - Relative and absolute tolerances. The solver keeps the local error - estimates less than ``atol + rtol * abs(y)``. Here `rtol` controls a - relative accuracy (number of correct digits), while `atol` controls - absolute accuracy (number of correct decimal places). To achieve the - desired `rtol`, set `atol` to be smaller than the smallest value that - can be expected from ``rtol * abs(y)`` so that `rtol` dominates the - allowable error. If `atol` is larger than ``rtol * abs(y)`` the - number of correct digits is not guaranteed. Conversely, to achieve the - desired `atol` set `rtol` such that ``rtol * abs(y)`` is always smaller - than `atol`. If components of y have different scales, it might be - beneficial to set different `atol` values for different components by - passing array_like with shape (n,) for `atol`. Default values are - 1e-3 for `rtol` and 1e-6 for `atol`. - vectorized : bool, optional - Whether `fun` is implemented in a vectorized fashion. Default is False. - - Attributes - ---------- - n : int - Number of equations. - status : string - Current status of the solver: 'running', 'finished' or 'failed'. - t_bound : float - Boundary time. - direction : float - Integration direction: +1 or -1. - t : float - Current time. - y : ndarray - Current state. - t_old : float - Previous time. None if no steps were made yet. - step_size : float - Size of the last successful step. None if no steps were made yet. - nfev : int - Number evaluations of the system's right-hand side. - njev : int - Number of evaluations of the Jacobian. Is always 0 for this solver - as it does not use the Jacobian. - nlu : int - Number of LU decompositions. Is always 0 for this solver. - - References - ---------- - .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential - Equations I: Nonstiff Problems", Sec. II. - .. [2] `Page with original Fortran code of DOP853 - `_. - """ - - n_stages = dop853_coefficients.N_STAGES - order = 8 - error_estimator_order = 7 - A = dop853_coefficients.A[:n_stages, :n_stages] - B = dop853_coefficients.B - C = dop853_coefficients.C[:n_stages] - E3 = dop853_coefficients.E3 - E5 = dop853_coefficients.E5 - D = dop853_coefficients.D - - A_EXTRA = dop853_coefficients.A[n_stages + 1 :] - C_EXTRA = dop853_coefficients.C[n_stages + 1 :] - - def __init__( - self, - fun, - t0, - y0, - t_bound, - max_step=np.inf, - rtol=1e-3, - atol=1e-6, - vectorized=False, - first_step=None, - **extraneous, - ): - super().__init__( - fun, - t0, - y0, - t_bound, - max_step, - rtol, - atol, - vectorized, - first_step, - **extraneous, - ) - self.K_extended = np.empty( - (dop853_coefficients.N_STAGES_EXTENDED, self.n), dtype=self.y.dtype - ) - self.K = self.K_extended[: self.n_stages + 1] - - def _estimate_error(self, K, h): # Left for testing purposes. - err5 = np.dot(K.T, self.E5) - err3 = np.dot(K.T, self.E3) - denom = np.hypot(np.abs(err5), 0.1 * np.abs(err3)) - correction_factor = np.ones_like(err5) - mask = denom > 0 - correction_factor[mask] = np.abs(err5[mask]) / denom[mask] - return h * err5 * correction_factor - - def _estimate_error_norm(self, K, h, scale): - err5 = np.dot(K.T, self.E5) / scale - err3 = np.dot(K.T, self.E3) / scale - err5_norm_2 = np.linalg.norm(err5) ** 2 - err3_norm_2 = np.linalg.norm(err3) ** 2 - if err5_norm_2 == 0 and err3_norm_2 == 0: - return 0.0 - denom = err5_norm_2 + 0.01 * err3_norm_2 - return np.abs(h) * err5_norm_2 / np.sqrt(denom * len(scale)) - def _dense_output_impl(self): K = self.K_extended h = self.h_previous @@ -754,3 +623,86 @@ def _dense_output_impl(self): F[3:] = h * np.dot(self.D, K) return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) + + +_ = """Explicit Runge-Kutta method of order 8. + +This is a Python implementation of "DOP853" algorithm originally written +in Fortran [1]_, [2]_. Note that this is not a literate translation, but +the algorithmic core and coefficients are the same. + +Can be applied in the complex domain. + +Parameters +---------- +fun : callable + Right-hand side of the system. The calling signature is ``fun(t, y)``. + Here, ``t`` is a scalar, and there are two options for the ndarray ``y``: + It can either have shape (n,); then ``fun`` must return array_like with + shape (n,). Alternatively it can have shape (n, k); then ``fun`` + must return an array_like with shape (n, k), i.e. each column + corresponds to a single column in ``y``. The choice between the two + options is determined by `vectorized` argument (see below). +t0 : float + Initial time. +y0 : array_like, shape (n,) + Initial state. +t_bound : float + Boundary time - the integration won't continue beyond it. It also + determines the direction of the integration. +first_step : float or None, optional + Initial step size. Default is ``None`` which means that the algorithm + should choose. +max_step : float, optional + Maximum allowed step size. Default is np.inf, i.e. the step size is not + bounded and determined solely by the solver. +rtol, atol : float and array_like, optional + Relative and absolute tolerances. The solver keeps the local error + estimates less than ``atol + rtol * abs(y)``. Here `rtol` controls a + relative accuracy (number of correct digits), while `atol` controls + absolute accuracy (number of correct decimal places). To achieve the + desired `rtol`, set `atol` to be smaller than the smallest value that + can be expected from ``rtol * abs(y)`` so that `rtol` dominates the + allowable error. If `atol` is larger than ``rtol * abs(y)`` the + number of correct digits is not guaranteed. Conversely, to achieve the + desired `atol` set `rtol` such that ``rtol * abs(y)`` is always smaller + than `atol`. If components of y have different scales, it might be + beneficial to set different `atol` values for different components by + passing array_like with shape (n,) for `atol`. Default values are + 1e-3 for `rtol` and 1e-6 for `atol`. +vectorized : bool, optional + Whether `fun` is implemented in a vectorized fashion. Default is False. + +Attributes +---------- +n : int + Number of equations. +status : string + Current status of the solver: 'running', 'finished' or 'failed'. +t_bound : float + Boundary time. +direction : float + Integration direction: +1 or -1. +t : float + Current time. +y : ndarray + Current state. +t_old : float + Previous time. None if no steps were made yet. +step_size : float + Size of the last successful step. None if no steps were made yet. +nfev : int + Number evaluations of the system's right-hand side. +njev : int + Number of evaluations of the Jacobian. Is always 0 for this solver + as it does not use the Jacobian. +nlu : int + Number of LU decompositions. Is always 0 for this solver. + +References +---------- +.. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential + Equations I: Nonstiff Problems", Sec. II. +.. [2] `Page with original Fortran code of DOP853 + `_. +""" From 9aec2486884554ebe9ea9ab295da2af235f0c527 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 13 Jan 2024 09:53:40 +0100 Subject: [PATCH 100/346] rm class --- src/hapsira/core/math/ivp/_rk.py | 124 ++++++++++++++----------------- 1 file changed, 56 insertions(+), 68 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 3c12cc829..f9e50cf27 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -268,7 +268,7 @@ def fun_wrapped(t, y): return fun_wrapped, y0 -class OdeSolver: +class DOP853: """Base class for ODE solvers. In order to implement a new solver you need to follow the guidelines: @@ -368,7 +368,35 @@ class OdeSolver: TOO_SMALL_STEP = "Required step size is less than spacing between numbers." - def __init__(self, fun, t0, y0, t_bound, vectorized): + n_stages: int = dop853_coefficients.N_STAGES + order: int = 8 + error_estimator_order: int = 7 + A = dop853_coefficients.A[:n_stages, :n_stages] + B = dop853_coefficients.B + C = dop853_coefficients.C[:n_stages] + E: np.ndarray = NotImplemented + E3 = dop853_coefficients.E3 + E5 = dop853_coefficients.E5 + D = dop853_coefficients.D + P: np.ndarray = NotImplemented + + A_EXTRA = dop853_coefficients.A[n_stages + 1 :] + C_EXTRA = dop853_coefficients.C[n_stages + 1 :] + + def __init__( + self, + fun, + t0, + y0, + t_bound, + max_step=np.inf, + rtol=1e-3, + atol=1e-6, + vectorized=False, + first_step=None, + **extraneous, + ): + warn_extraneous(extraneous) self.t_old = None self.t = t0 self._fun, self.y = check_arguments(fun, y0) @@ -406,6 +434,32 @@ def fun(t, y): self.njev = 0 self.nlu = 0 + self.y_old = None + self.max_step = validate_max_step(max_step) + self.rtol, self.atol = validate_tol(rtol, atol, self.n) + self.f = self.fun(self.t, self.y) + if first_step is None: + self.h_abs = select_initial_step( + self.fun, + self.t, + self.y, + self.f, + self.direction, + self.error_estimator_order, + self.rtol, + self.atol, + ) + else: + self.h_abs = validate_first_step(first_step, t0, t_bound) + self.K = np.empty((self.n_stages + 1, self.n), dtype=self.y.dtype) + self.error_exponent = -1 / (self.error_estimator_order + 1) + self.h_previous = None + + self.K_extended = np.empty( + (dop853_coefficients.N_STAGES_EXTENDED, self.n), dtype=self.y.dtype + ) + self.K = self.K_extended[: self.n_stages + 1] + @property def step_size(self): if self.t_old is None: @@ -459,72 +513,6 @@ def dense_output(self): return self._dense_output_impl() - def _step_impl(self): - raise NotImplementedError - - def _dense_output_impl(self): - raise NotImplementedError - - -class DOP853(OdeSolver): - """Base class for explicit Runge-Kutta methods.""" - - n_stages: int = dop853_coefficients.N_STAGES - order: int = 8 - error_estimator_order: int = 7 - A = dop853_coefficients.A[:n_stages, :n_stages] - B = dop853_coefficients.B - C = dop853_coefficients.C[:n_stages] - E: np.ndarray = NotImplemented - E3 = dop853_coefficients.E3 - E5 = dop853_coefficients.E5 - D = dop853_coefficients.D - P: np.ndarray = NotImplemented - - A_EXTRA = dop853_coefficients.A[n_stages + 1 :] - C_EXTRA = dop853_coefficients.C[n_stages + 1 :] - - def __init__( - self, - fun, - t0, - y0, - t_bound, - max_step=np.inf, - rtol=1e-3, - atol=1e-6, - vectorized=False, - first_step=None, - **extraneous, - ): - warn_extraneous(extraneous) - super().__init__(fun, t0, y0, t_bound, vectorized) - self.y_old = None - self.max_step = validate_max_step(max_step) - self.rtol, self.atol = validate_tol(rtol, atol, self.n) - self.f = self.fun(self.t, self.y) - if first_step is None: - self.h_abs = select_initial_step( - self.fun, - self.t, - self.y, - self.f, - self.direction, - self.error_estimator_order, - self.rtol, - self.atol, - ) - else: - self.h_abs = validate_first_step(first_step, t0, t_bound) - self.K = np.empty((self.n_stages + 1, self.n), dtype=self.y.dtype) - self.error_exponent = -1 / (self.error_estimator_order + 1) - self.h_previous = None - - self.K_extended = np.empty( - (dop853_coefficients.N_STAGES_EXTENDED, self.n), dtype=self.y.dtype - ) - self.K = self.K_extended[: self.n_stages + 1] - def _estimate_error_norm(self, K, h, scale): err5 = np.dot(K.T, self.E5) / scale err3 = np.dot(K.T, self.E3) / scale From bd4b8cba37f68fc4d0b7c2e51b49981acecca492 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 13 Jan 2024 09:58:02 +0100 Subject: [PATCH 101/346] cleanup --- src/hapsira/core/math/ivp/_rk.py | 185 +++++++------------------------ 1 file changed, 39 insertions(+), 146 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index f9e50cf27..bf593ebe9 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -269,76 +269,51 @@ def fun_wrapped(t, y): class DOP853: - """Base class for ODE solvers. - - In order to implement a new solver you need to follow the guidelines: - - 1. A constructor must accept parameters presented in the base class - (listed below) along with any other parameters specific to a solver. - 2. A constructor must accept arbitrary extraneous arguments - ``**extraneous``, but warn that these arguments are irrelevant - using `common.warn_extraneous` function. Do not pass these - arguments to the base class. - 3. A solver must implement a private method `_step_impl(self)` which - propagates a solver one step further. It must return tuple - ``(success, message)``, where ``success`` is a boolean indicating - whether a step was successful, and ``message`` is a string - containing description of a failure if a step failed or None - otherwise. - 4. A solver must implement a private method `_dense_output_impl(self)`, - which returns a `DenseOutput` object covering the last successful - step. - 5. A solver must have attributes listed below in Attributes section. - Note that ``t_old`` and ``step_size`` are updated automatically. - 6. Use `fun(self, t, y)` method for the system rhs evaluation, this - way the number of function evaluations (`nfev`) will be tracked - automatically. - 7. For convenience, a base class provides `fun_single(self, t, y)` and - `fun_vectorized(self, t, y)` for evaluating the rhs in - non-vectorized and vectorized fashions respectively (regardless of - how `fun` from the constructor is implemented). These calls don't - increment `nfev`. - 8. If a solver uses a Jacobian matrix and LU decompositions, it should - track the number of Jacobian evaluations (`njev`) and the number of - LU decompositions (`nlu`). - 9. By convention, the function evaluations used to compute a finite - difference approximation of the Jacobian should not be counted in - `nfev`, thus use `fun_single(self, t, y)` or - `fun_vectorized(self, t, y)` when computing a finite difference - approximation of the Jacobian. + """Explicit Runge-Kutta method of order 8. + + This is a Python implementation of "DOP853" algorithm originally written + in Fortran [1]_, [2]_. Note that this is not a literate translation, but + the algorithmic core and coefficients are the same. Parameters ---------- fun : callable - Right-hand side of the system: the time derivative of the state ``y`` - at time ``t``. The calling signature is ``fun(t, y)``, where ``t`` is a - scalar and ``y`` is an ndarray with ``len(y) = len(y0)``. ``fun`` must - return an array of the same shape as ``y``. See `vectorized` for more - information. + Right-hand side of the system. The calling signature is ``fun(t, y)``. + Here, ``t`` is a scalar, and there are two options for the ndarray ``y``: + It can either have shape (n,); then ``fun`` must return array_like with + shape (n,). Alternatively it can have shape (n, k); then ``fun`` + must return an array_like with shape (n, k), i.e. each column + corresponds to a single column in ``y``. The choice between the two + options is determined by `vectorized` argument (see below). t0 : float Initial time. y0 : array_like, shape (n,) Initial state. t_bound : float - Boundary time --- the integration won't continue beyond it. It also + Boundary time - the integration won't continue beyond it. It also determines the direction of the integration. - vectorized : bool - Whether `fun` can be called in a vectorized fashion. Default is False. - - If ``vectorized`` is False, `fun` will always be called with ``y`` of - shape ``(n,)``, where ``n = len(y0)``. - - If ``vectorized`` is True, `fun` may be called with ``y`` of shape - ``(n, k)``, where ``k`` is an integer. In this case, `fun` must behave - such that ``fun(t, y)[:, i] == fun(t, y[:, i])`` (i.e. each column of - the returned array is the time derivative of the state corresponding - with a column of ``y``). - - Setting ``vectorized=True`` allows for faster finite difference - approximation of the Jacobian by methods 'Radau' and 'BDF', but - will result in slower execution for other methods. It can also - result in slower overall execution for 'Radau' and 'BDF' in some - circumstances (e.g. small ``len(y0)``). + first_step : float or None, optional + Initial step size. Default is ``None`` which means that the algorithm + should choose. + max_step : float, optional + Maximum allowed step size. Default is np.inf, i.e. the step size is not + bounded and determined solely by the solver. + rtol, atol : float and array_like, optional + Relative and absolute tolerances. The solver keeps the local error + estimates less than ``atol + rtol * abs(y)``. Here `rtol` controls a + relative accuracy (number of correct digits), while `atol` controls + absolute accuracy (number of correct decimal places). To achieve the + desired `rtol`, set `atol` to be smaller than the smallest value that + can be expected from ``rtol * abs(y)`` so that `rtol` dominates the + allowable error. If `atol` is larger than ``rtol * abs(y)`` the + number of correct digits is not guaranteed. Conversely, to achieve the + desired `atol` set `rtol` such that ``rtol * abs(y)`` is always smaller + than `atol`. If components of y have different scales, it might be + beneficial to set different `atol` values for different components by + passing array_like with shape (n,) for `atol`. Default values are + 1e-3 for `rtol` and 1e-6 for `atol`. + vectorized : bool, optional + Whether `fun` is implemented in a vectorized fashion. Default is False. Attributes ---------- @@ -359,11 +334,12 @@ class DOP853: step_size : float Size of the last successful step. None if no steps were made yet. nfev : int - Number of the system's rhs evaluations. + Number evaluations of the system's right-hand side. njev : int - Number of the Jacobian evaluations. + Number of evaluations of the Jacobian. Is always 0 for this solver + as it does not use the Jacobian. nlu : int - Number of LU decompositions. + Number of LU decompositions. Is always 0 for this solver. """ TOO_SMALL_STEP = "Required step size is less than spacing between numbers." @@ -611,86 +587,3 @@ def _dense_output_impl(self): F[3:] = h * np.dot(self.D, K) return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) - - -_ = """Explicit Runge-Kutta method of order 8. - -This is a Python implementation of "DOP853" algorithm originally written -in Fortran [1]_, [2]_. Note that this is not a literate translation, but -the algorithmic core and coefficients are the same. - -Can be applied in the complex domain. - -Parameters ----------- -fun : callable - Right-hand side of the system. The calling signature is ``fun(t, y)``. - Here, ``t`` is a scalar, and there are two options for the ndarray ``y``: - It can either have shape (n,); then ``fun`` must return array_like with - shape (n,). Alternatively it can have shape (n, k); then ``fun`` - must return an array_like with shape (n, k), i.e. each column - corresponds to a single column in ``y``. The choice between the two - options is determined by `vectorized` argument (see below). -t0 : float - Initial time. -y0 : array_like, shape (n,) - Initial state. -t_bound : float - Boundary time - the integration won't continue beyond it. It also - determines the direction of the integration. -first_step : float or None, optional - Initial step size. Default is ``None`` which means that the algorithm - should choose. -max_step : float, optional - Maximum allowed step size. Default is np.inf, i.e. the step size is not - bounded and determined solely by the solver. -rtol, atol : float and array_like, optional - Relative and absolute tolerances. The solver keeps the local error - estimates less than ``atol + rtol * abs(y)``. Here `rtol` controls a - relative accuracy (number of correct digits), while `atol` controls - absolute accuracy (number of correct decimal places). To achieve the - desired `rtol`, set `atol` to be smaller than the smallest value that - can be expected from ``rtol * abs(y)`` so that `rtol` dominates the - allowable error. If `atol` is larger than ``rtol * abs(y)`` the - number of correct digits is not guaranteed. Conversely, to achieve the - desired `atol` set `rtol` such that ``rtol * abs(y)`` is always smaller - than `atol`. If components of y have different scales, it might be - beneficial to set different `atol` values for different components by - passing array_like with shape (n,) for `atol`. Default values are - 1e-3 for `rtol` and 1e-6 for `atol`. -vectorized : bool, optional - Whether `fun` is implemented in a vectorized fashion. Default is False. - -Attributes ----------- -n : int - Number of equations. -status : string - Current status of the solver: 'running', 'finished' or 'failed'. -t_bound : float - Boundary time. -direction : float - Integration direction: +1 or -1. -t : float - Current time. -y : ndarray - Current state. -t_old : float - Previous time. None if no steps were made yet. -step_size : float - Size of the last successful step. None if no steps were made yet. -nfev : int - Number evaluations of the system's right-hand side. -njev : int - Number of evaluations of the Jacobian. Is always 0 for this solver - as it does not use the Jacobian. -nlu : int - Number of LU decompositions. Is always 0 for this solver. - -References ----------- -.. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential - Equations I: Nonstiff Problems", Sec. II. -.. [2] `Page with original Fortran code of DOP853 - `_. -""" From 3546e037be5c981fe1f07281c6456d3bb0e3c9ac Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 13 Jan 2024 09:59:09 +0100 Subject: [PATCH 102/346] exports --- src/hapsira/core/math/ivp/_rk.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index bf593ebe9..d4f075bf5 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -5,6 +5,12 @@ from . import _dop853_coefficients as dop853_coefficients +__all__ = [ + "EPS", + "DOP853", +] + + EPS = np.finfo(float).eps # Multiply steps computed from asymptotic behaviour of errors by this. From c585852a5d88f1eb481650b697faff2e85cf1823 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 14 Jan 2024 16:10:51 +0100 Subject: [PATCH 103/346] rm arg --- src/hapsira/core/propagation/cowell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index c0959cc15..0324fa87e 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -18,7 +18,7 @@ def cowell(k, r, v, tofs, rtol=1e-11, *, events=None, f=func_twobody): args=(k,), rtol=rtol, atol=1e-12, - dense_output=True, + # dense_output=True, events=events, ) if not result.success: From 0bbb45334764215333dfba13982e5f111b5a1c8f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 14 Jan 2024 16:11:18 +0100 Subject: [PATCH 104/346] types, const, cleanup --- src/hapsira/core/math/ivp/_rk.py | 73 +++++++++++++++++--------------- 1 file changed, 38 insertions(+), 35 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index d4f075bf5..f9b632ecd 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -1,6 +1,8 @@ +from typing import Callable, Tuple from warnings import warn import numpy as np +from numba import jit from . import _dop853_coefficients as dop853_coefficients @@ -20,12 +22,24 @@ MAX_FACTOR = 10 # Maximum allowed increase in a step size. -def norm(x): +@jit(nopython=False) +def norm(x: np.ndarray) -> float: """Compute RMS norm.""" return np.linalg.norm(x) / x.size**0.5 -def rk_step(fun, t, y, f, h, A, B, C, K): +@jit(nopython=False) +def rk_step( + fun: Callable, + t: float, + y: np.ndarray, + f: np.ndarray, + h: float, + A: np.ndarray, + B: np.ndarray, + C: np.ndarray, + K: np.ndarray, +) -> Tuple[np.ndarray, np.ndarray]: """Perform a single Runge-Kutta step. This function computes a prediction of an explicit Runge-Kutta method and @@ -85,7 +99,17 @@ def rk_step(fun, t, y, f, h, A, B, C, K): return y_new, f_new -def select_initial_step(fun, t0, y0, f0, direction, order, rtol, atol): +@jit(nopython=False) +def select_initial_step( + fun: Callable, + t0: float, + y0: np.ndarray, + f0: np.ndarray, + direction: float, + order: float, + rtol: float, + atol: float, +) -> float: """Empirically select a good initial step. The algorithm is described in [1]_. @@ -143,7 +167,8 @@ def select_initial_step(fun, t0, y0, f0, direction, order, rtol, atol): return min(100 * h0, h1) -def validate_first_step(first_step, t0, t_bound): +@jit(nopython=False) +def validate_first_step(first_step: float, t0: float, t_bound: float) -> float: """Assert that first_step is valid and return it.""" if first_step <= 0: raise ValueError("`first_step` must be positive.") @@ -152,34 +177,16 @@ def validate_first_step(first_step, t0, t_bound): return first_step -def validate_max_step(max_step): +@jit(nopython=False) +def validate_max_step(max_step: float) -> float: """Assert that max_Step is valid and return it.""" if max_step <= 0: raise ValueError("`max_step` must be positive.") return max_step -def warn_extraneous(extraneous): - """Display a warning for extraneous keyword arguments. - - The initializer of each solver class is expected to collect keyword - arguments that it doesn't understand and warn about them. This function - prints a warning for each key in the supplied dictionary. - - Parameters - ---------- - extraneous : dict - Extraneous keyword arguments - """ - if extraneous: - warn( - "The following arguments have no effect for a chosen solver: {}.".format( - ", ".join(f"`{x}`" for x in extraneous) - ) - ) - - -def validate_tol(rtol, atol, n): +@jit(nopython=False) +def validate_tol(rtol: float, atol: float, n: int) -> Tuple[float, float]: """Validate tolerance values.""" if np.any(rtol < 100 * EPS): @@ -350,7 +357,7 @@ class DOP853: TOO_SMALL_STEP = "Required step size is less than spacing between numbers." - n_stages: int = dop853_coefficients.N_STAGES + n_stages: int = 12 # N_STAGES == 12 order: int = 8 error_estimator_order: int = 7 A = dop853_coefficients.A[:n_stages, :n_stages] @@ -368,7 +375,7 @@ class DOP853: def __init__( self, fun, - t0, + t0: float, y0, t_bound, max_step=np.inf, @@ -376,9 +383,7 @@ def __init__( atol=1e-6, vectorized=False, first_step=None, - **extraneous, ): - warn_extraneous(extraneous) self.t_old = None self.t = t0 self._fun, self.y = check_arguments(fun, y0) @@ -438,8 +443,8 @@ def fun(t, y): self.h_previous = None self.K_extended = np.empty( - (dop853_coefficients.N_STAGES_EXTENDED, self.n), dtype=self.y.dtype - ) + (16, self.n), dtype=self.y.dtype + ) # N_STAGES_EXTENDED == 16 self.K = self.K_extended[: self.n_stages + 1] @property @@ -580,9 +585,7 @@ def _dense_output_impl(self): dy = np.dot(K[:s].T, a[:s]) * h K[s] = self.fun(self.t_old + c * h, self.y_old + dy) - F = np.empty( - (dop853_coefficients.INTERPOLATOR_POWER, self.n), dtype=self.y_old.dtype - ) + F = np.empty((7, self.n), dtype=self.y_old.dtype) # INTERPOLATOR_POWER==7 f_old = K[0] delta_y = self.y - self.y_old From 57a963c59e0b1157741bd08755650cd7e5a7e302 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 14 Jan 2024 16:28:37 +0100 Subject: [PATCH 105/346] revert solver --- contrib/CR3BP/CR3BP.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/contrib/CR3BP/CR3BP.py b/contrib/CR3BP/CR3BP.py index d581e4683..2a78611d1 100644 --- a/contrib/CR3BP/CR3BP.py +++ b/contrib/CR3BP/CR3BP.py @@ -32,7 +32,8 @@ from numba import njit as jit import numpy as np -from hapsira.core.math.ivp import solve_ivp +# from hapsira.core.math.ivp import solve_ivp +from scipy.integrate import solve_ivp, DOP853 @jit @@ -314,6 +315,7 @@ def propagate(mu, r0, v0, tofs, rtol=1e-11, f=func_CR3BP): args=(mu,), rtol=rtol, atol=1e-12, + method=DOP853, dense_output=True, ) @@ -370,6 +372,7 @@ def propagateSTM(mu, r0, v0, STM0, tofs, rtol=1e-11, f=func_CR3BP_STM): args=(mu,), rtol=rtol, atol=1e-12, + method=DOP853, dense_output=True, ) From f43d4c9698a26bf082e9cc99b70eaea34d5407a1 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 14 Jan 2024 17:11:16 +0100 Subject: [PATCH 106/346] rm vectorized support --- src/hapsira/core/math/ivp/_api.py | 19 +----------------- src/hapsira/core/math/ivp/_rk.py | 32 +++++++------------------------ 2 files changed, 8 insertions(+), 43 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 4f0afd6e9..30eaad464 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -178,7 +178,6 @@ def solve_ivp( # t_eval=None, # dense_output=False, events=None, - vectorized=False, args=None, **options, ): @@ -230,22 +229,6 @@ def solve_ivp( You can assign attributes like ``event.terminal = True`` to any function in Python. - vectorized : bool, optional - Whether `fun` can be called in a vectorized fashion. Default is False. - - If ``vectorized`` is False, `fun` will always be called with ``y`` of - shape ``(n,)``, where ``n = len(y0)``. - - If ``vectorized`` is True, `fun` may be called with ``y`` of shape - ``(n, k)``, where ``k`` is an integer. In this case, `fun` must behave - such that ``fun(t, y)[:, i] == fun(t, y[:, i])`` (i.e. each column of - the returned array is the time derivative of the state corresponding - with a column of ``y``). - - Setting ``vectorized=True`` allows for faster finite difference - approximation of the Jacobian by methods 'Radau' and 'BDF', but - will result in slower execution for other methods and for 'Radau' and - 'BDF' in some circumstances (e.g. small ``len(y0)``). args : tuple, optional Additional arguments to pass to the user-defined functions. If given, the additional arguments are passed to all user-defined functions. @@ -359,7 +342,7 @@ def solve_ivp( def fun(t, x, fun=fun): return fun(t, x, *args) - solver = method(fun, t0, y0, tf, vectorized=vectorized, **options) + solver = method(fun, t0, y0, tf, **options) ts = [t0] ys = [y0] diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index f9b632ecd..f2ef88153 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -325,8 +325,6 @@ class DOP853: beneficial to set different `atol` values for different components by passing array_like with shape (n,) for `atol`. Default values are 1e-3 for `rtol` and 1e-6 for `atol`. - vectorized : bool, optional - Whether `fun` is implemented in a vectorized fashion. Default is False. Attributes ---------- @@ -374,36 +372,21 @@ class DOP853: def __init__( self, - fun, + fun: Callable, t0: float, - y0, - t_bound, - max_step=np.inf, - rtol=1e-3, - atol=1e-6, - vectorized=False, + y0: np.array, + t_bound: float, + max_step: float = np.inf, + rtol: float = 1e-3, + atol: float = 1e-6, first_step=None, ): self.t_old = None self.t = t0 self._fun, self.y = check_arguments(fun, y0) self.t_bound = t_bound - self.vectorized = vectorized - if vectorized: - - def fun_single(t, y): - return self._fun(t, y[:, None]).ravel() - - fun_vectorized = self._fun - else: - fun_single = self._fun - - def fun_vectorized(t, y): - f = np.empty_like(y) - for i, yi in enumerate(y.T): - f[:, i] = self._fun(t, yi) - return f + fun_single = self._fun def fun(t, y): self.nfev += 1 @@ -411,7 +394,6 @@ def fun(t, y): self.fun = fun self.fun_single = fun_single - self.fun_vectorized = fun_vectorized self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 self.n = self.y.size From ee77d0736314cc651ac02d5dfa98c1cbe0a3fa46 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 14 Jan 2024 18:03:05 +0100 Subject: [PATCH 107/346] rm first step --- src/hapsira/core/math/ivp/_rk.py | 37 +++++++++----------------------- 1 file changed, 10 insertions(+), 27 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index f2ef88153..5cd77250b 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -167,16 +167,6 @@ def select_initial_step( return min(100 * h0, h1) -@jit(nopython=False) -def validate_first_step(first_step: float, t0: float, t_bound: float) -> float: - """Assert that first_step is valid and return it.""" - if first_step <= 0: - raise ValueError("`first_step` must be positive.") - if first_step > np.abs(t_bound - t0): - raise ValueError("`first_step` exceeds bounds.") - return first_step - - @jit(nopython=False) def validate_max_step(max_step: float) -> float: """Assert that max_Step is valid and return it.""" @@ -305,9 +295,6 @@ class DOP853: t_bound : float Boundary time - the integration won't continue beyond it. It also determines the direction of the integration. - first_step : float or None, optional - Initial step size. Default is ``None`` which means that the algorithm - should choose. max_step : float, optional Maximum allowed step size. Default is np.inf, i.e. the step size is not bounded and determined solely by the solver. @@ -379,7 +366,6 @@ def __init__( max_step: float = np.inf, rtol: float = 1e-3, atol: float = 1e-6, - first_step=None, ): self.t_old = None self.t = t0 @@ -407,19 +393,16 @@ def fun(t, y): self.max_step = validate_max_step(max_step) self.rtol, self.atol = validate_tol(rtol, atol, self.n) self.f = self.fun(self.t, self.y) - if first_step is None: - self.h_abs = select_initial_step( - self.fun, - self.t, - self.y, - self.f, - self.direction, - self.error_estimator_order, - self.rtol, - self.atol, - ) - else: - self.h_abs = validate_first_step(first_step, t0, t_bound) + self.h_abs = select_initial_step( + self.fun, + self.t, + self.y, + self.f, + self.direction, + self.error_estimator_order, + self.rtol, + self.atol, + ) self.K = np.empty((self.n_stages + 1, self.n), dtype=self.y.dtype) self.error_exponent = -1 / (self.error_estimator_order + 1) self.h_previous = None From 81d5611900b70d3b858afb103f6b283f2f2d2fb2 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 14 Jan 2024 18:04:20 +0100 Subject: [PATCH 108/346] rm oo k decl --- src/hapsira/core/math/ivp/_rk.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 5cd77250b..012867042 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -403,7 +403,6 @@ def fun(t, y): self.rtol, self.atol, ) - self.K = np.empty((self.n_stages + 1, self.n), dtype=self.y.dtype) self.error_exponent = -1 / (self.error_estimator_order + 1) self.h_previous = None From 203357e3e36d2f3d421ab29cbd06f0832a91d14a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 14 Jan 2024 18:15:56 +0100 Subject: [PATCH 109/346] fixed array sizes, prep for tuples --- src/hapsira/core/math/ivp/_rk.py | 39 ++++++++++++++++---------------- 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 012867042..6c257a4d3 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -21,6 +21,11 @@ MIN_FACTOR = 0.2 # Minimum allowed decrease in a step size. MAX_FACTOR = 10 # Maximum allowed increase in a step size. +INTERPOLATOR_POWER = 7 +N_RV = 6 +N_STAGES = 12 +N_STAGES_EXTENDED = 16 + @jit(nopython=False) def norm(x: np.ndarray) -> float: @@ -176,7 +181,7 @@ def validate_max_step(max_step: float) -> float: @jit(nopython=False) -def validate_tol(rtol: float, atol: float, n: int) -> Tuple[float, float]: +def validate_tol(rtol: float, atol: float) -> Tuple[float, float]: """Validate tolerance values.""" if np.any(rtol < 100 * EPS): @@ -187,7 +192,7 @@ def validate_tol(rtol: float, atol: float, n: int) -> Tuple[float, float]: rtol = np.maximum(rtol, 100 * EPS) atol = np.asarray(atol) - if atol.ndim > 0 and atol.shape != (n,): + if atol.ndim > 0 and atol.shape != (N_RV,): raise ValueError("`atol` has wrong shape.") if np.any(atol < 0): @@ -342,20 +347,19 @@ class DOP853: TOO_SMALL_STEP = "Required step size is less than spacing between numbers." - n_stages: int = 12 # N_STAGES == 12 order: int = 8 error_estimator_order: int = 7 - A = dop853_coefficients.A[:n_stages, :n_stages] + A = dop853_coefficients.A[:N_STAGES, :N_STAGES] B = dop853_coefficients.B - C = dop853_coefficients.C[:n_stages] + C = dop853_coefficients.C[:N_STAGES] E: np.ndarray = NotImplemented E3 = dop853_coefficients.E3 E5 = dop853_coefficients.E5 D = dop853_coefficients.D P: np.ndarray = NotImplemented - A_EXTRA = dop853_coefficients.A[n_stages + 1 :] - C_EXTRA = dop853_coefficients.C[n_stages + 1 :] + A_EXTRA = dop853_coefficients.A[N_STAGES + 1 :] + C_EXTRA = dop853_coefficients.C[N_STAGES + 1 :] def __init__( self, @@ -367,6 +371,8 @@ def __init__( rtol: float = 1e-3, atol: float = 1e-6, ): + assert y0.shape == (N_RV,) + self.t_old = None self.t = t0 self._fun, self.y = check_arguments(fun, y0) @@ -382,7 +388,6 @@ def fun(t, y): self.fun_single = fun_single self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 - self.n = self.y.size self.status = "running" self.nfev = 0 @@ -391,7 +396,7 @@ def fun(t, y): self.y_old = None self.max_step = validate_max_step(max_step) - self.rtol, self.atol = validate_tol(rtol, atol, self.n) + self.rtol, self.atol = validate_tol(rtol, atol) self.f = self.fun(self.t, self.y) self.h_abs = select_initial_step( self.fun, @@ -406,10 +411,8 @@ def fun(t, y): self.error_exponent = -1 / (self.error_estimator_order + 1) self.h_previous = None - self.K_extended = np.empty( - (16, self.n), dtype=self.y.dtype - ) # N_STAGES_EXTENDED == 16 - self.K = self.K_extended[: self.n_stages + 1] + self.K_extended = np.empty((N_STAGES_EXTENDED, N_RV), dtype=self.y.dtype) + self.K = self.K_extended[: N_STAGES + 1] @property def step_size(self): @@ -431,7 +434,7 @@ def step(self): if self.status != "running": raise RuntimeError("Attempt to step on a failed or finished " "solver.") - if self.n == 0 or self.t == self.t_bound: + if self.t == self.t_bound: # Handle corner cases of empty solver or no integration. self.t_old = self.t self.t = self.t_bound @@ -460,7 +463,7 @@ def dense_output(self): """ assert self.t_old is not None - assert not (self.n == 0 or self.t == self.t_old) + assert self.t != self.t_old return self._dense_output_impl() @@ -543,13 +546,11 @@ def _step_impl(self): def _dense_output_impl(self): K = self.K_extended h = self.h_previous - for s, (a, c) in enumerate( - zip(self.A_EXTRA, self.C_EXTRA), start=self.n_stages + 1 - ): + for s, (a, c) in enumerate(zip(self.A_EXTRA, self.C_EXTRA), start=N_STAGES + 1): dy = np.dot(K[:s].T, a[:s]) * h K[s] = self.fun(self.t_old + c * h, self.y_old + dy) - F = np.empty((7, self.n), dtype=self.y_old.dtype) # INTERPOLATOR_POWER==7 + F = np.empty((INTERPOLATOR_POWER, N_RV), dtype=self.y_old.dtype) f_old = K[0] delta_y = self.y - self.y_old From 0c548af1ae36a0a24fc95857b815203205f14f34 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 15 Jan 2024 12:08:01 +0100 Subject: [PATCH 110/346] nb compile cache clean --- Makefile | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Makefile b/Makefile index 0fa102438..bc75d96f4 100644 --- a/Makefile +++ b/Makefile @@ -5,6 +5,8 @@ _clean_coverage: coverage erase _clean_py: + find src/ tests/ contrib/ -name '*.nbi' -exec rm -f {} + + find src/ tests/ contrib/ -name '*.nbc' -exec rm -f {} + find src/ tests/ contrib/ -name '*.pyc' -exec rm -f {} + find src/ tests/ contrib/ -name '*.pyo' -exec rm -f {} + find src/ tests/ contrib/ -name '*~' -exec rm -f {} + From 7cbfeaecccb9f40f60aca7ab12f0f488b679b31b Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 15 Jan 2024 12:08:36 +0100 Subject: [PATCH 111/346] explicit nopython --- src/hapsira/core/math/ivp/_brentq.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index b9d6993e6..f27de4dc9 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -27,7 +27,7 @@ def _signbit_hf(a): return a < 0 -@jit +@jit(nopython=False) def _brentq_hf( func, # callback_type xa, # double @@ -107,7 +107,7 @@ def _brentq_hf( return xcur, funcalls, iterations, CONVERR -@jit +@jit(nopython=False) def brentq_sf( func, # func a, # double From 7c5defc055f0bf6bb6310d7f30c9b74144b6180f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 15 Jan 2024 12:08:56 +0100 Subject: [PATCH 112/346] assert shapes; unroll loop --- src/hapsira/core/math/ivp/_rk.py | 51 ++++++++++++++++++++++++++++++-- 1 file changed, 48 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 6c257a4d3..f6caf9ce8 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -91,16 +91,61 @@ def rk_step( .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential Equations I: Nonstiff Problems", Sec. II.4. """ + + assert y.shape == (N_RV,) + assert f.shape == (N_RV,) + assert A.shape == (N_STAGES, N_STAGES) + assert B.shape == (N_STAGES,) + assert C.shape == (N_STAGES,) + assert K.shape == (N_STAGES + 1, N_RV) + K[0] = f - for s, (a, c) in enumerate(zip(A[1:], C[1:]), start=1): - dy = np.dot(K[:s].T, a[:s]) * h - K[s] = fun(t + c * h, y + dy) + + # for s, (a, c) in enumerate(zip(A[1:], C[1:]), start=1): + # dy = np.dot(K[:s].T, a[:s]) * h + # K[s] = fun(t + c * h, y + dy) + + dy = np.dot(K[:1].T, A[1, :1]) * h + K[1] = fun(t + C[1] * h, y + dy) + + dy = np.dot(K[:2].T, A[2, :2]) * h + K[2] = fun(t + C[2] * h, y + dy) + + dy = np.dot(K[:3].T, A[3, :3]) * h + K[3] = fun(t + C[3] * h, y + dy) + + dy = np.dot(K[:4].T, A[4, :4]) * h + K[4] = fun(t + C[4] * h, y + dy) + + dy = np.dot(K[:5].T, A[5, :5]) * h + K[5] = fun(t + C[5] * h, y + dy) + + dy = np.dot(K[:6].T, A[6, :6]) * h + K[6] = fun(t + C[6] * h, y + dy) + + dy = np.dot(K[:7].T, A[7, :7]) * h + K[7] = fun(t + C[7] * h, y + dy) + + dy = np.dot(K[:8].T, A[8, :8]) * h + K[8] = fun(t + C[8] * h, y + dy) + + dy = np.dot(K[:9].T, A[9, :9]) * h + K[9] = fun(t + C[9] * h, y + dy) + + dy = np.dot(K[:10].T, A[10, :10]) * h + K[10] = fun(t + C[10] * h, y + dy) + + dy = np.dot(K[:11].T, A[11, :11]) * h + K[11] = fun(t + C[11] * h, y + dy) y_new = y + h * np.dot(K[:-1].T, B) f_new = fun(t + h, y_new) K[-1] = f_new + assert y_new.shape == (N_RV,) + assert f_new.shape == (N_RV,) + return y_new, f_new From adb7ca6db51d149c78c295c01cbfa4300c7abf63 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 15 Jan 2024 12:45:17 +0100 Subject: [PATCH 113/346] rm A,B,C props and args, turn into const --- src/hapsira/core/math/ivp/_rk.py | 42 ++++++++++++++++---------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index f6caf9ce8..586928e17 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -5,7 +5,7 @@ from numba import jit from . import _dop853_coefficients as dop853_coefficients - +from ._dop853_coefficients import A, B, C __all__ = [ "EPS", @@ -26,6 +26,10 @@ N_STAGES = 12 N_STAGES_EXTENDED = 16 +A = A[:N_STAGES, :N_STAGES] +B = B +C = C[:N_STAGES] + @jit(nopython=False) def norm(x: np.ndarray) -> float: @@ -40,9 +44,6 @@ def rk_step( y: np.ndarray, f: np.ndarray, h: float, - A: np.ndarray, - B: np.ndarray, - C: np.ndarray, K: np.ndarray, ) -> Tuple[np.ndarray, np.ndarray]: """Perform a single Runge-Kutta step. @@ -64,16 +65,6 @@ def rk_step( Current value of the derivative, i.e., ``fun(x, y)``. h : float Step to use. - A : ndarray, shape (n_stages, n_stages) - Coefficients for combining previous RK stages to compute the next - stage. For explicit methods the coefficients at and above the main - diagonal are zeros. - B : ndarray, shape (n_stages,) - Coefficients for combining RK stages for computing the final - prediction. - C : ndarray, shape (n_stages,) - Coefficients for incrementing time for consecutive RK stages. - The value for the first stage is always zero. K : ndarray, shape (n_stages + 1, n) Storage array for putting RK stages here. Stages are stored in rows. The last row is a linear combination of the previous rows with @@ -86,6 +77,19 @@ def rk_step( f_new : ndarray, shape (n,) Derivative ``fun(t + h, y_new)``. + Const + ----- + A : ndarray, shape (n_stages, n_stages) + Coefficients for combining previous RK stages to compute the next + stage. For explicit methods the coefficients at and above the main + diagonal are zeros. + B : ndarray, shape (n_stages,) + Coefficients for combining RK stages for computing the final + prediction. + C : ndarray, shape (n_stages,) + Coefficients for incrementing time for consecutive RK stages. + The value for the first stage is always zero. + References ---------- .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential @@ -94,10 +98,11 @@ def rk_step( assert y.shape == (N_RV,) assert f.shape == (N_RV,) + assert K.shape == (N_STAGES + 1, N_RV) + assert A.shape == (N_STAGES, N_STAGES) assert B.shape == (N_STAGES,) assert C.shape == (N_STAGES,) - assert K.shape == (N_STAGES + 1, N_RV) K[0] = f @@ -394,9 +399,6 @@ class DOP853: order: int = 8 error_estimator_order: int = 7 - A = dop853_coefficients.A[:N_STAGES, :N_STAGES] - B = dop853_coefficients.B - C = dop853_coefficients.C[:N_STAGES] E: np.ndarray = NotImplemented E3 = dop853_coefficients.E3 E5 = dop853_coefficients.E5 @@ -555,9 +557,7 @@ def _step_impl(self): h = t_new - t h_abs = np.abs(h) - y_new, f_new = rk_step( - self.fun, t, y, self.f, h, self.A, self.B, self.C, self.K - ) + y_new, f_new = rk_step(self.fun, t, y, self.f, h, self.K) scale = atol + np.maximum(np.abs(y), np.abs(y_new)) * rtol error_norm = self._estimate_error_norm(self.K, h, scale) From bd49a33efd17a986648150e056b0fd732f36368c Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 15 Jan 2024 13:22:39 +0100 Subject: [PATCH 114/346] eliminate func wrapper --- src/hapsira/core/math/ivp/_api.py | 9 --------- src/hapsira/core/math/ivp/_rk.py | 21 +++------------------ 2 files changed, 3 insertions(+), 27 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 30eaad464..c3313fdfb 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -24,9 +24,6 @@ class OdeResult(dict): underlying solver. Refer to `message` for details. message : str Description of the cause of the termination. - nfev, njev : int - Number of evaluations of the objective functions and of its - Jacobian and Hessian. t, y, sol, t_events, y_events, nlu : ? """ @@ -314,10 +311,6 @@ def solve_ivp( y_events : list of ndarray or None For each value of `t_events`, the corresponding value of the solution. None if `events` was None. - nfev : int - Number of evaluations of the right-hand side. - njev : int - Number of evaluations of the Jacobian. nlu : int Number of LU decompositions. status : int @@ -422,8 +415,6 @@ def fun(t, x, fun=fun): sol=sol, t_events=t_events, y_events=y_events, - nfev=solver.nfev, - njev=solver.njev, nlu=solver.nlu, status=status, message=MESSAGES.get(status, message), diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 586928e17..59204a43d 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -307,7 +307,7 @@ def _call_impl(self, t): return y.T -def check_arguments(fun, y0): +def check_arguments(y0): """Helper function for checking arguments common to all solvers.""" y0 = np.asarray(y0) @@ -320,10 +320,7 @@ def check_arguments(fun, y0): assert not y0.ndim != 1 assert np.isfinite(y0).all() - def fun_wrapped(t, y): - return np.asarray(fun(t, y), dtype=dtype) - - return fun_wrapped, y0 + return y0 class DOP853: @@ -386,11 +383,6 @@ class DOP853: Previous time. None if no steps were made yet. step_size : float Size of the last successful step. None if no steps were made yet. - nfev : int - Number evaluations of the system's right-hand side. - njev : int - Number of evaluations of the Jacobian. Is always 0 for this solver - as it does not use the Jacobian. nlu : int Number of LU decompositions. Is always 0 for this solver. """ @@ -422,17 +414,10 @@ def __init__( self.t_old = None self.t = t0 - self._fun, self.y = check_arguments(fun, y0) + self.y = check_arguments(y0) self.t_bound = t_bound - fun_single = self._fun - - def fun(t, y): - self.nfev += 1 - return self.fun_single(t, y) - self.fun = fun - self.fun_single = fun_single self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 self.status = "running" From 0322830b996d3e3365fab79a2447a6c0cc64dae9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 15 Jan 2024 13:43:44 +0100 Subject: [PATCH 115/346] pass k to func, prep for inline compile and eliminate last wrapper --- src/hapsira/core/math/ivp/_api.py | 22 +++------------ src/hapsira/core/math/ivp/_rk.py | 37 +++++++++++++++----------- src/hapsira/core/propagation/cowell.py | 2 +- 3 files changed, 25 insertions(+), 36 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index c3313fdfb..348ead74f 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -171,11 +171,11 @@ def solve_ivp( fun, t_span, y0, + argk: float, method=DOP853, # t_eval=None, # dense_output=False, events=None, - args=None, **options, ): """Solve an initial value problem for a system of ODEs. @@ -226,12 +226,6 @@ def solve_ivp( You can assign attributes like ``event.terminal = True`` to any function in Python. - args : tuple, optional - Additional arguments to pass to the user-defined functions. If given, - the additional arguments are passed to all user-defined functions. - So if, for example, `fun` has the signature ``fun(t, y, a, b, c)``, - then `jac` (if given) and any event functions must have the same - signature, and `args` must be a tuple of length 3. **options Options passed to a chosen solver. All options available for already implemented solvers are listed below. @@ -330,12 +324,7 @@ def solve_ivp( t0, tf = map(float, t_span) - assert isinstance(args, tuple) - - def fun(t, x, fun=fun): - return fun(t, x, *args) - - solver = method(fun, t0, y0, tf, **options) + solver = method(fun, t0, y0, tf, argk, **options) ts = [t0] ys = [y0] @@ -345,12 +334,7 @@ def fun(t, x, fun=fun): events, is_terminal, event_dir = prepare_events(events) if events is not None: - if args is not None: - # Wrap user functions in lambdas to hide the additional parameters. - # The original event function is passed as a keyword argument to the - # lambda to keep the original function in scope (i.e., avoid the - # late binding closure "gotcha"). - events = [lambda t, x, event=event: event(t, x, *args) for event in events] + events = [lambda t, x, event=event: event(t, x, argk) for event in events] g = [event(t0, y0) for event in events] t_events = [[] for _ in range(len(events))] y_events = [[] for _ in range(len(events))] diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 59204a43d..53819ea9e 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -45,6 +45,7 @@ def rk_step( f: np.ndarray, h: float, K: np.ndarray, + argk: float, ) -> Tuple[np.ndarray, np.ndarray]: """Perform a single Runge-Kutta step. @@ -111,40 +112,40 @@ def rk_step( # K[s] = fun(t + c * h, y + dy) dy = np.dot(K[:1].T, A[1, :1]) * h - K[1] = fun(t + C[1] * h, y + dy) + K[1] = fun(t + C[1] * h, y + dy, argk) dy = np.dot(K[:2].T, A[2, :2]) * h - K[2] = fun(t + C[2] * h, y + dy) + K[2] = fun(t + C[2] * h, y + dy, argk) dy = np.dot(K[:3].T, A[3, :3]) * h - K[3] = fun(t + C[3] * h, y + dy) + K[3] = fun(t + C[3] * h, y + dy, argk) dy = np.dot(K[:4].T, A[4, :4]) * h - K[4] = fun(t + C[4] * h, y + dy) + K[4] = fun(t + C[4] * h, y + dy, argk) dy = np.dot(K[:5].T, A[5, :5]) * h - K[5] = fun(t + C[5] * h, y + dy) + K[5] = fun(t + C[5] * h, y + dy, argk) dy = np.dot(K[:6].T, A[6, :6]) * h - K[6] = fun(t + C[6] * h, y + dy) + K[6] = fun(t + C[6] * h, y + dy, argk) dy = np.dot(K[:7].T, A[7, :7]) * h - K[7] = fun(t + C[7] * h, y + dy) + K[7] = fun(t + C[7] * h, y + dy, argk) dy = np.dot(K[:8].T, A[8, :8]) * h - K[8] = fun(t + C[8] * h, y + dy) + K[8] = fun(t + C[8] * h, y + dy, argk) dy = np.dot(K[:9].T, A[9, :9]) * h - K[9] = fun(t + C[9] * h, y + dy) + K[9] = fun(t + C[9] * h, y + dy, argk) dy = np.dot(K[:10].T, A[10, :10]) * h - K[10] = fun(t + C[10] * h, y + dy) + K[10] = fun(t + C[10] * h, y + dy, argk) dy = np.dot(K[:11].T, A[11, :11]) * h - K[11] = fun(t + C[11] * h, y + dy) + K[11] = fun(t + C[11] * h, y + dy, argk) y_new = y + h * np.dot(K[:-1].T, B) - f_new = fun(t + h, y_new) + f_new = fun(t + h, y_new, argk) K[-1] = f_new @@ -159,6 +160,7 @@ def select_initial_step( fun: Callable, t0: float, y0: np.ndarray, + argk: float, f0: np.ndarray, direction: float, order: float, @@ -211,7 +213,7 @@ def select_initial_step( h0 = 0.01 * d0 / d1 y1 = y0 + h0 * direction * f0 - f1 = fun(t0 + h0 * direction, y1) + f1 = fun(t0 + h0 * direction, y1, argk) d2 = norm((f1 - f0) / scale) / h0 if d1 <= 1e-15 and d2 <= 1e-15: @@ -406,6 +408,7 @@ def __init__( t0: float, y0: np.array, t_bound: float, + argk: float, max_step: float = np.inf, rtol: float = 1e-3, atol: float = 1e-6, @@ -418,6 +421,7 @@ def __init__( self.t_bound = t_bound self.fun = fun + self.argk = argk self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 self.status = "running" @@ -429,11 +433,12 @@ def __init__( self.y_old = None self.max_step = validate_max_step(max_step) self.rtol, self.atol = validate_tol(rtol, atol) - self.f = self.fun(self.t, self.y) + self.f = self.fun(self.t, self.y, self.argk) self.h_abs = select_initial_step( self.fun, self.t, self.y, + self.argk, self.f, self.direction, self.error_estimator_order, @@ -542,7 +547,7 @@ def _step_impl(self): h = t_new - t h_abs = np.abs(h) - y_new, f_new = rk_step(self.fun, t, y, self.f, h, self.K) + y_new, f_new = rk_step(self.fun, t, y, self.f, h, self.K, self.argk) scale = atol + np.maximum(np.abs(y), np.abs(y_new)) * rtol error_norm = self._estimate_error_norm(self.K, h, scale) @@ -578,7 +583,7 @@ def _dense_output_impl(self): h = self.h_previous for s, (a, c) in enumerate(zip(self.A_EXTRA, self.C_EXTRA), start=N_STAGES + 1): dy = np.dot(K[:s].T, a[:s]) * h - K[s] = self.fun(self.t_old + c * h, self.y_old + dy) + K[s] = self.fun(self.t_old + c * h, self.y_old + dy, self.argk) F = np.empty((INTERPOLATOR_POWER, N_RV), dtype=self.y_old.dtype) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 0324fa87e..a28d33024 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -15,7 +15,7 @@ def cowell(k, r, v, tofs, rtol=1e-11, *, events=None, f=func_twobody): f, (0, max(tofs)), u0, - args=(k,), + argk=k, rtol=rtol, atol=1e-12, # dense_output=True, From d2225c8cdb5e3325a964610fcaa9268600dd7ce6 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 15 Jan 2024 14:01:41 +0100 Subject: [PATCH 116/346] rm dead consts --- src/hapsira/core/math/ivp/_rk.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 53819ea9e..c80c5a8a2 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -25,6 +25,7 @@ N_RV = 6 N_STAGES = 12 N_STAGES_EXTENDED = 16 +ERROR_ESTIMATOR_ORDER = 7 A = A[:N_STAGES, :N_STAGES] B = B @@ -391,13 +392,9 @@ class DOP853: TOO_SMALL_STEP = "Required step size is less than spacing between numbers." - order: int = 8 - error_estimator_order: int = 7 - E: np.ndarray = NotImplemented E3 = dop853_coefficients.E3 E5 = dop853_coefficients.E5 D = dop853_coefficients.D - P: np.ndarray = NotImplemented A_EXTRA = dop853_coefficients.A[N_STAGES + 1 :] C_EXTRA = dop853_coefficients.C[N_STAGES + 1 :] @@ -441,11 +438,11 @@ def __init__( self.argk, self.f, self.direction, - self.error_estimator_order, + ERROR_ESTIMATOR_ORDER, self.rtol, self.atol, ) - self.error_exponent = -1 / (self.error_estimator_order + 1) + self.error_exponent = -1 / (ERROR_ESTIMATOR_ORDER + 1) self.h_previous = None self.K_extended = np.empty((N_STAGES_EXTENDED, N_RV), dtype=self.y.dtype) From bceaa1a019285f3079e895b1cbd5076ceabe0cba Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 15 Jan 2024 16:06:07 +0100 Subject: [PATCH 117/346] prep for tuples without numpy in rk step --- src/hapsira/core/math/ivp/_rk.py | 588 ++++++++++++++++++++++++++++++- 1 file changed, 577 insertions(+), 11 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index c80c5a8a2..746536696 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -112,37 +112,603 @@ def rk_step( # dy = np.dot(K[:s].T, a[:s]) * h # K[s] = fun(t + c * h, y + dy) - dy = np.dot(K[:1].T, A[1, :1]) * h + # dy = np.dot(K[:1].T, A[1, :1]) * h + dy = np.array( + [ + (K[0, 0] * A[1, 0]) * h, + (K[0, 1] * A[1, 0]) * h, + (K[0, 2] * A[1, 0]) * h, + (K[0, 3] * A[1, 0]) * h, + (K[0, 4] * A[1, 0]) * h, + (K[0, 5] * A[1, 0]) * h, + ] + ) K[1] = fun(t + C[1] * h, y + dy, argk) - dy = np.dot(K[:2].T, A[2, :2]) * h + # dy = np.dot(K[:2].T, A[2, :2]) * h + dy = np.array( + [ + (K[0, 0] * A[2, 0] + K[1, 0] * A[2, 1]) * h, + (K[0, 1] * A[2, 0] + K[1, 1] * A[2, 1]) * h, + (K[0, 2] * A[2, 0] + K[1, 2] * A[2, 1]) * h, + (K[0, 3] * A[2, 0] + K[1, 3] * A[2, 1]) * h, + (K[0, 4] * A[2, 0] + K[1, 4] * A[2, 1]) * h, + (K[0, 5] * A[2, 0] + K[1, 5] * A[2, 1]) * h, + ] + ) K[2] = fun(t + C[2] * h, y + dy, argk) - dy = np.dot(K[:3].T, A[3, :3]) * h + # dy = np.dot(K[:3].T, A[3, :3]) * h + dy = np.array( + [ + (K[0, 0] * A[3, 0] + K[1, 0] * A[3, 1] + K[2, 0] * A[3, 2]) * h, + (K[0, 1] * A[3, 0] + K[1, 1] * A[3, 1] + K[2, 1] * A[3, 2]) * h, + (K[0, 2] * A[3, 0] + K[1, 2] * A[3, 1] + K[2, 2] * A[3, 2]) * h, + (K[0, 3] * A[3, 0] + K[1, 3] * A[3, 1] + K[2, 3] * A[3, 2]) * h, + (K[0, 4] * A[3, 0] + K[1, 4] * A[3, 1] + K[2, 4] * A[3, 2]) * h, + (K[0, 5] * A[3, 0] + K[1, 5] * A[3, 1] + K[2, 5] * A[3, 2]) * h, + ] + ) K[3] = fun(t + C[3] * h, y + dy, argk) - dy = np.dot(K[:4].T, A[4, :4]) * h + # dy = np.dot(K[:4].T, A[4, :4]) * h + dy = np.array( + [ + ( + K[0, 0] * A[4, 0] + + K[1, 0] * A[4, 1] + + K[2, 0] * A[4, 2] + + K[3, 0] * A[4, 3] + ) + * h, + ( + K[0, 1] * A[4, 0] + + K[1, 1] * A[4, 1] + + K[2, 1] * A[4, 2] + + K[3, 1] * A[4, 3] + ) + * h, + ( + K[0, 2] * A[4, 0] + + K[1, 2] * A[4, 1] + + K[2, 2] * A[4, 2] + + K[3, 2] * A[4, 3] + ) + * h, + ( + K[0, 3] * A[4, 0] + + K[1, 3] * A[4, 1] + + K[2, 3] * A[4, 2] + + K[3, 3] * A[4, 3] + ) + * h, + ( + K[0, 4] * A[4, 0] + + K[1, 4] * A[4, 1] + + K[2, 4] * A[4, 2] + + K[3, 4] * A[4, 3] + ) + * h, + ( + K[0, 5] * A[4, 0] + + K[1, 5] * A[4, 1] + + K[2, 5] * A[4, 2] + + K[3, 5] * A[4, 3] + ) + * h, + ] + ) K[4] = fun(t + C[4] * h, y + dy, argk) - dy = np.dot(K[:5].T, A[5, :5]) * h + # dy = np.dot(K[:5].T, A[5, :5]) * h + dy = np.array( + [ + ( + K[0, 0] * A[5, 0] + + K[1, 0] * A[5, 1] + + K[2, 0] * A[5, 2] + + K[3, 0] * A[5, 3] + + K[4, 0] * A[5, 4] + ) + * h, + ( + K[0, 1] * A[5, 0] + + K[1, 1] * A[5, 1] + + K[2, 1] * A[5, 2] + + K[3, 1] * A[5, 3] + + K[4, 1] * A[5, 4] + ) + * h, + ( + K[0, 2] * A[5, 0] + + K[1, 2] * A[5, 1] + + K[2, 2] * A[5, 2] + + K[3, 2] * A[5, 3] + + K[4, 2] * A[5, 4] + ) + * h, + ( + K[0, 3] * A[5, 0] + + K[1, 3] * A[5, 1] + + K[2, 3] * A[5, 2] + + K[3, 3] * A[5, 3] + + K[4, 3] * A[5, 4] + ) + * h, + ( + K[0, 4] * A[5, 0] + + K[1, 4] * A[5, 1] + + K[2, 4] * A[5, 2] + + K[3, 4] * A[5, 3] + + K[4, 4] * A[5, 4] + ) + * h, + ( + K[0, 5] * A[5, 0] + + K[1, 5] * A[5, 1] + + K[2, 5] * A[5, 2] + + K[3, 5] * A[5, 3] + + K[4, 5] * A[5, 4] + ) + * h, + ] + ) K[5] = fun(t + C[5] * h, y + dy, argk) - dy = np.dot(K[:6].T, A[6, :6]) * h + # dy = np.dot(K[:6].T, A[6, :6]) * h + dy = np.array( + [ + ( + K[0, 0] * A[6, 0] + + K[1, 0] * A[6, 1] + + K[2, 0] * A[6, 2] + + K[3, 0] * A[6, 3] + + K[4, 0] * A[6, 4] + + K[5, 0] * A[6, 5] + ) + * h, + ( + K[0, 1] * A[6, 0] + + K[1, 1] * A[6, 1] + + K[2, 1] * A[6, 2] + + K[3, 1] * A[6, 3] + + K[4, 1] * A[6, 4] + + K[5, 1] * A[6, 5] + ) + * h, + ( + K[0, 2] * A[6, 0] + + K[1, 2] * A[6, 1] + + K[2, 2] * A[6, 2] + + K[3, 2] * A[6, 3] + + K[4, 2] * A[6, 4] + + K[5, 2] * A[6, 5] + ) + * h, + ( + K[0, 3] * A[6, 0] + + K[1, 3] * A[6, 1] + + K[2, 3] * A[6, 2] + + K[3, 3] * A[6, 3] + + K[4, 3] * A[6, 4] + + K[5, 3] * A[6, 5] + ) + * h, + ( + K[0, 4] * A[6, 0] + + K[1, 4] * A[6, 1] + + K[2, 4] * A[6, 2] + + K[3, 4] * A[6, 3] + + K[4, 4] * A[6, 4] + + K[5, 4] * A[6, 5] + ) + * h, + ( + K[0, 5] * A[6, 0] + + K[1, 5] * A[6, 1] + + K[2, 5] * A[6, 2] + + K[3, 5] * A[6, 3] + + K[4, 5] * A[6, 4] + + K[5, 5] * A[6, 5] + ) + * h, + ] + ) K[6] = fun(t + C[6] * h, y + dy, argk) - dy = np.dot(K[:7].T, A[7, :7]) * h + # dy = np.dot(K[:7].T, A[7, :7]) * h + dy = np.array( + [ + ( + K[0, 0] * A[7, 0] + + K[1, 0] * A[7, 1] + + K[2, 0] * A[7, 2] + + K[3, 0] * A[7, 3] + + K[4, 0] * A[7, 4] + + K[5, 0] * A[7, 5] + + K[6, 0] * A[7, 6] + ) + * h, + ( + K[0, 1] * A[7, 0] + + K[1, 1] * A[7, 1] + + K[2, 1] * A[7, 2] + + K[3, 1] * A[7, 3] + + K[4, 1] * A[7, 4] + + K[5, 1] * A[7, 5] + + K[6, 1] * A[7, 6] + ) + * h, + ( + K[0, 2] * A[7, 0] + + K[1, 2] * A[7, 1] + + K[2, 2] * A[7, 2] + + K[3, 2] * A[7, 3] + + K[4, 2] * A[7, 4] + + K[5, 2] * A[7, 5] + + K[6, 2] * A[7, 6] + ) + * h, + ( + K[0, 3] * A[7, 0] + + K[1, 3] * A[7, 1] + + K[2, 3] * A[7, 2] + + K[3, 3] * A[7, 3] + + K[4, 3] * A[7, 4] + + K[5, 3] * A[7, 5] + + K[6, 3] * A[7, 6] + ) + * h, + ( + K[0, 4] * A[7, 0] + + K[1, 4] * A[7, 1] + + K[2, 4] * A[7, 2] + + K[3, 4] * A[7, 3] + + K[4, 4] * A[7, 4] + + K[5, 4] * A[7, 5] + + K[6, 4] * A[7, 6] + ) + * h, + ( + K[0, 5] * A[7, 0] + + K[1, 5] * A[7, 1] + + K[2, 5] * A[7, 2] + + K[3, 5] * A[7, 3] + + K[4, 5] * A[7, 4] + + K[5, 5] * A[7, 5] + + K[6, 5] * A[7, 6] + ) + * h, + ] + ) K[7] = fun(t + C[7] * h, y + dy, argk) - dy = np.dot(K[:8].T, A[8, :8]) * h + # dy = np.dot(K[:8].T, A[8, :8]) * h + dy = np.array( + [ + ( + K[0, 0] * A[8, 0] + + K[1, 0] * A[8, 1] + + K[2, 0] * A[8, 2] + + K[3, 0] * A[8, 3] + + K[4, 0] * A[8, 4] + + K[5, 0] * A[8, 5] + + K[6, 0] * A[8, 6] + + K[7, 0] * A[8, 7] + ) + * h, + ( + K[0, 1] * A[8, 0] + + K[1, 1] * A[8, 1] + + K[2, 1] * A[8, 2] + + K[3, 1] * A[8, 3] + + K[4, 1] * A[8, 4] + + K[5, 1] * A[8, 5] + + K[6, 1] * A[8, 6] + + K[7, 1] * A[8, 7] + ) + * h, + ( + K[0, 2] * A[8, 0] + + K[1, 2] * A[8, 1] + + K[2, 2] * A[8, 2] + + K[3, 2] * A[8, 3] + + K[4, 2] * A[8, 4] + + K[5, 2] * A[8, 5] + + K[6, 2] * A[8, 6] + + K[7, 2] * A[8, 7] + ) + * h, + ( + K[0, 3] * A[8, 0] + + K[1, 3] * A[8, 1] + + K[2, 3] * A[8, 2] + + K[3, 3] * A[8, 3] + + K[4, 3] * A[8, 4] + + K[5, 3] * A[8, 5] + + K[6, 3] * A[8, 6] + + K[7, 3] * A[8, 7] + ) + * h, + ( + K[0, 4] * A[8, 0] + + K[1, 4] * A[8, 1] + + K[2, 4] * A[8, 2] + + K[3, 4] * A[8, 3] + + K[4, 4] * A[8, 4] + + K[5, 4] * A[8, 5] + + K[6, 4] * A[8, 6] + + K[7, 4] * A[8, 7] + ) + * h, + ( + K[0, 5] * A[8, 0] + + K[1, 5] * A[8, 1] + + K[2, 5] * A[8, 2] + + K[3, 5] * A[8, 3] + + K[4, 5] * A[8, 4] + + K[5, 5] * A[8, 5] + + K[6, 5] * A[8, 6] + + K[7, 5] * A[8, 7] + ) + * h, + ] + ) K[8] = fun(t + C[8] * h, y + dy, argk) - dy = np.dot(K[:9].T, A[9, :9]) * h + # dy = np.dot(K[:9].T, A[9, :9]) * h + dy = np.array( + [ + ( + K[0, 0] * A[9, 0] + + K[1, 0] * A[9, 1] + + K[2, 0] * A[9, 2] + + K[3, 0] * A[9, 3] + + K[4, 0] * A[9, 4] + + K[5, 0] * A[9, 5] + + K[6, 0] * A[9, 6] + + K[7, 0] * A[9, 7] + + K[8, 0] * A[9, 8] + ) + * h, + ( + K[0, 1] * A[9, 0] + + K[1, 1] * A[9, 1] + + K[2, 1] * A[9, 2] + + K[3, 1] * A[9, 3] + + K[4, 1] * A[9, 4] + + K[5, 1] * A[9, 5] + + K[6, 1] * A[9, 6] + + K[7, 1] * A[9, 7] + + K[8, 1] * A[9, 8] + ) + * h, + ( + K[0, 2] * A[9, 0] + + K[1, 2] * A[9, 1] + + K[2, 2] * A[9, 2] + + K[3, 2] * A[9, 3] + + K[4, 2] * A[9, 4] + + K[5, 2] * A[9, 5] + + K[6, 2] * A[9, 6] + + K[7, 2] * A[9, 7] + + K[8, 2] * A[9, 8] + ) + * h, + ( + K[0, 3] * A[9, 0] + + K[1, 3] * A[9, 1] + + K[2, 3] * A[9, 2] + + K[3, 3] * A[9, 3] + + K[4, 3] * A[9, 4] + + K[5, 3] * A[9, 5] + + K[6, 3] * A[9, 6] + + K[7, 3] * A[9, 7] + + K[8, 3] * A[9, 8] + ) + * h, + ( + K[0, 4] * A[9, 0] + + K[1, 4] * A[9, 1] + + K[2, 4] * A[9, 2] + + K[3, 4] * A[9, 3] + + K[4, 4] * A[9, 4] + + K[5, 4] * A[9, 5] + + K[6, 4] * A[9, 6] + + K[7, 4] * A[9, 7] + + K[8, 4] * A[9, 8] + ) + * h, + ( + K[0, 5] * A[9, 0] + + K[1, 5] * A[9, 1] + + K[2, 5] * A[9, 2] + + K[3, 5] * A[9, 3] + + K[4, 5] * A[9, 4] + + K[5, 5] * A[9, 5] + + K[6, 5] * A[9, 6] + + K[7, 5] * A[9, 7] + + K[8, 5] * A[9, 8] + ) + * h, + ] + ) K[9] = fun(t + C[9] * h, y + dy, argk) - dy = np.dot(K[:10].T, A[10, :10]) * h + # dy = np.dot(K[:10].T, A[10, :10]) * h + dy = np.array( + [ + ( + K[0, 0] * A[10, 0] + + K[1, 0] * A[10, 1] + + K[2, 0] * A[10, 2] + + K[3, 0] * A[10, 3] + + K[4, 0] * A[10, 4] + + K[5, 0] * A[10, 5] + + K[6, 0] * A[10, 6] + + K[7, 0] * A[10, 7] + + K[8, 0] * A[10, 8] + + K[9, 0] * A[10, 9] + ) + * h, + ( + K[0, 1] * A[10, 0] + + K[1, 1] * A[10, 1] + + K[2, 1] * A[10, 2] + + K[3, 1] * A[10, 3] + + K[4, 1] * A[10, 4] + + K[5, 1] * A[10, 5] + + K[6, 1] * A[10, 6] + + K[7, 1] * A[10, 7] + + K[8, 1] * A[10, 8] + + K[9, 1] * A[10, 9] + ) + * h, + ( + K[0, 2] * A[10, 0] + + K[1, 2] * A[10, 1] + + K[2, 2] * A[10, 2] + + K[3, 2] * A[10, 3] + + K[4, 2] * A[10, 4] + + K[5, 2] * A[10, 5] + + K[6, 2] * A[10, 6] + + K[7, 2] * A[10, 7] + + K[8, 2] * A[10, 8] + + K[9, 2] * A[10, 9] + ) + * h, + ( + K[0, 3] * A[10, 0] + + K[1, 3] * A[10, 1] + + K[2, 3] * A[10, 2] + + K[3, 3] * A[10, 3] + + K[4, 3] * A[10, 4] + + K[5, 3] * A[10, 5] + + K[6, 3] * A[10, 6] + + K[7, 3] * A[10, 7] + + K[8, 3] * A[10, 8] + + K[9, 3] * A[10, 9] + ) + * h, + ( + K[0, 4] * A[10, 0] + + K[1, 4] * A[10, 1] + + K[2, 4] * A[10, 2] + + K[3, 4] * A[10, 3] + + K[4, 4] * A[10, 4] + + K[5, 4] * A[10, 5] + + K[6, 4] * A[10, 6] + + K[7, 4] * A[10, 7] + + K[8, 4] * A[10, 8] + + K[9, 4] * A[10, 9] + ) + * h, + ( + K[0, 5] * A[10, 0] + + K[1, 5] * A[10, 1] + + K[2, 5] * A[10, 2] + + K[3, 5] * A[10, 3] + + K[4, 5] * A[10, 4] + + K[5, 5] * A[10, 5] + + K[6, 5] * A[10, 6] + + K[7, 5] * A[10, 7] + + K[8, 5] * A[10, 8] + + K[9, 5] * A[10, 9] + ) + * h, + ] + ) K[10] = fun(t + C[10] * h, y + dy, argk) - dy = np.dot(K[:11].T, A[11, :11]) * h + # dy = np.dot(K[:11].T, A[11, :11]) * h + dy = np.array( + [ + ( + K[0, 0] * A[11, 0] + + K[1, 0] * A[11, 1] + + K[2, 0] * A[11, 2] + + K[3, 0] * A[11, 3] + + K[4, 0] * A[11, 4] + + K[5, 0] * A[11, 5] + + K[6, 0] * A[11, 6] + + K[7, 0] * A[11, 7] + + K[8, 0] * A[11, 8] + + K[9, 0] * A[11, 9] + + K[10, 0] * A[11, 10] + ) + * h, + ( + K[0, 1] * A[11, 0] + + K[1, 1] * A[11, 1] + + K[2, 1] * A[11, 2] + + K[3, 1] * A[11, 3] + + K[4, 1] * A[11, 4] + + K[5, 1] * A[11, 5] + + K[6, 1] * A[11, 6] + + K[7, 1] * A[11, 7] + + K[8, 1] * A[11, 8] + + K[9, 1] * A[11, 9] + + K[10, 1] * A[11, 10] + ) + * h, + ( + K[0, 2] * A[11, 0] + + K[1, 2] * A[11, 1] + + K[2, 2] * A[11, 2] + + K[3, 2] * A[11, 3] + + K[4, 2] * A[11, 4] + + K[5, 2] * A[11, 5] + + K[6, 2] * A[11, 6] + + K[7, 2] * A[11, 7] + + K[8, 2] * A[11, 8] + + K[9, 2] * A[11, 9] + + K[10, 2] * A[11, 10] + ) + * h, + ( + K[0, 3] * A[11, 0] + + K[1, 3] * A[11, 1] + + K[2, 3] * A[11, 2] + + K[3, 3] * A[11, 3] + + K[4, 3] * A[11, 4] + + K[5, 3] * A[11, 5] + + K[6, 3] * A[11, 6] + + K[7, 3] * A[11, 7] + + K[8, 3] * A[11, 8] + + K[9, 3] * A[11, 9] + + K[10, 3] * A[11, 10] + ) + * h, + ( + K[0, 4] * A[11, 0] + + K[1, 4] * A[11, 1] + + K[2, 4] * A[11, 2] + + K[3, 4] * A[11, 3] + + K[4, 4] * A[11, 4] + + K[5, 4] * A[11, 5] + + K[6, 4] * A[11, 6] + + K[7, 4] * A[11, 7] + + K[8, 4] * A[11, 8] + + K[9, 4] * A[11, 9] + + K[10, 4] * A[11, 10] + ) + * h, + ( + K[0, 5] * A[11, 0] + + K[1, 5] * A[11, 1] + + K[2, 5] * A[11, 2] + + K[3, 5] * A[11, 3] + + K[4, 5] * A[11, 4] + + K[5, 5] * A[11, 5] + + K[6, 5] * A[11, 6] + + K[7, 5] * A[11, 7] + + K[8, 5] * A[11, 8] + + K[9, 5] * A[11, 9] + + K[10, 5] * A[11, 10] + ) + * h, + ] + ) K[11] = fun(t + C[11] * h, y + dy, argk) y_new = y + h * np.dot(K[:-1].T, B) From 10266a2d023367d193eb2391294dad48510be651 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 15 Jan 2024 16:23:58 +0100 Subject: [PATCH 118/346] prep for tuples without numpy in rk step (2) --- src/hapsira/core/math/ivp/_rk.py | 758 ++++++++++++++++--------------- 1 file changed, 380 insertions(+), 378 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 746536696..6dd47dff2 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -115,12 +115,12 @@ def rk_step( # dy = np.dot(K[:1].T, A[1, :1]) * h dy = np.array( [ - (K[0, 0] * A[1, 0]) * h, - (K[0, 1] * A[1, 0]) * h, - (K[0, 2] * A[1, 0]) * h, - (K[0, 3] * A[1, 0]) * h, - (K[0, 4] * A[1, 0]) * h, - (K[0, 5] * A[1, 0]) * h, + (K[0][0] * A[1][0]) * h, + (K[0][1] * A[1][0]) * h, + (K[0][2] * A[1][0]) * h, + (K[0][3] * A[1][0]) * h, + (K[0][4] * A[1][0]) * h, + (K[0][5] * A[1][0]) * h, ] ) K[1] = fun(t + C[1] * h, y + dy, argk) @@ -128,12 +128,12 @@ def rk_step( # dy = np.dot(K[:2].T, A[2, :2]) * h dy = np.array( [ - (K[0, 0] * A[2, 0] + K[1, 0] * A[2, 1]) * h, - (K[0, 1] * A[2, 0] + K[1, 1] * A[2, 1]) * h, - (K[0, 2] * A[2, 0] + K[1, 2] * A[2, 1]) * h, - (K[0, 3] * A[2, 0] + K[1, 3] * A[2, 1]) * h, - (K[0, 4] * A[2, 0] + K[1, 4] * A[2, 1]) * h, - (K[0, 5] * A[2, 0] + K[1, 5] * A[2, 1]) * h, + (K[0][0] * A[2][0] + K[1][0] * A[2][1]) * h, + (K[0][1] * A[2][0] + K[1][1] * A[2][1]) * h, + (K[0][2] * A[2][0] + K[1][2] * A[2][1]) * h, + (K[0][3] * A[2][0] + K[1][3] * A[2][1]) * h, + (K[0][4] * A[2][0] + K[1][4] * A[2][1]) * h, + (K[0][5] * A[2][0] + K[1][5] * A[2][1]) * h, ] ) K[2] = fun(t + C[2] * h, y + dy, argk) @@ -141,59 +141,60 @@ def rk_step( # dy = np.dot(K[:3].T, A[3, :3]) * h dy = np.array( [ - (K[0, 0] * A[3, 0] + K[1, 0] * A[3, 1] + K[2, 0] * A[3, 2]) * h, - (K[0, 1] * A[3, 0] + K[1, 1] * A[3, 1] + K[2, 1] * A[3, 2]) * h, - (K[0, 2] * A[3, 0] + K[1, 2] * A[3, 1] + K[2, 2] * A[3, 2]) * h, - (K[0, 3] * A[3, 0] + K[1, 3] * A[3, 1] + K[2, 3] * A[3, 2]) * h, - (K[0, 4] * A[3, 0] + K[1, 4] * A[3, 1] + K[2, 4] * A[3, 2]) * h, - (K[0, 5] * A[3, 0] + K[1, 5] * A[3, 1] + K[2, 5] * A[3, 2]) * h, + (K[0][0] * A[3][0] + K[1][0] * A[3][1] + K[2][0] * A[3][2]) * h, + (K[0][1] * A[3][0] + K[1][1] * A[3][1] + K[2][1] * A[3][2]) * h, + (K[0][2] * A[3][0] + K[1][2] * A[3][1] + K[2][2] * A[3][2]) * h, + (K[0][3] * A[3][0] + K[1][3] * A[3][1] + K[2][3] * A[3][2]) * h, + (K[0][4] * A[3][0] + K[1][4] * A[3][1] + K[2][4] * A[3][2]) * h, + (K[0][5] * A[3][0] + K[1][5] * A[3][1] + K[2][5] * A[3][2]) * h, ] ) + K[3] = fun(t + C[3] * h, y + dy, argk) # dy = np.dot(K[:4].T, A[4, :4]) * h dy = np.array( [ ( - K[0, 0] * A[4, 0] - + K[1, 0] * A[4, 1] - + K[2, 0] * A[4, 2] - + K[3, 0] * A[4, 3] + K[0][0] * A[4][0] + + K[1][0] * A[4][1] + + K[2][0] * A[4][2] + + K[3][0] * A[4][3] ) * h, ( - K[0, 1] * A[4, 0] - + K[1, 1] * A[4, 1] - + K[2, 1] * A[4, 2] - + K[3, 1] * A[4, 3] + K[0][1] * A[4][0] + + K[1][1] * A[4][1] + + K[2][1] * A[4][2] + + K[3][1] * A[4][3] ) * h, ( - K[0, 2] * A[4, 0] - + K[1, 2] * A[4, 1] - + K[2, 2] * A[4, 2] - + K[3, 2] * A[4, 3] + K[0][2] * A[4][0] + + K[1][2] * A[4][1] + + K[2][2] * A[4][2] + + K[3][2] * A[4][3] ) * h, ( - K[0, 3] * A[4, 0] - + K[1, 3] * A[4, 1] - + K[2, 3] * A[4, 2] - + K[3, 3] * A[4, 3] + K[0][3] * A[4][0] + + K[1][3] * A[4][1] + + K[2][3] * A[4][2] + + K[3][3] * A[4][3] ) * h, ( - K[0, 4] * A[4, 0] - + K[1, 4] * A[4, 1] - + K[2, 4] * A[4, 2] - + K[3, 4] * A[4, 3] + K[0][4] * A[4][0] + + K[1][4] * A[4][1] + + K[2][4] * A[4][2] + + K[3][4] * A[4][3] ) * h, ( - K[0, 5] * A[4, 0] - + K[1, 5] * A[4, 1] - + K[2, 5] * A[4, 2] - + K[3, 5] * A[4, 3] + K[0][5] * A[4][0] + + K[1][5] * A[4][1] + + K[2][5] * A[4][2] + + K[3][5] * A[4][3] ) * h, ] @@ -204,51 +205,51 @@ def rk_step( dy = np.array( [ ( - K[0, 0] * A[5, 0] - + K[1, 0] * A[5, 1] - + K[2, 0] * A[5, 2] - + K[3, 0] * A[5, 3] - + K[4, 0] * A[5, 4] + K[0][0] * A[5][0] + + K[1][0] * A[5][1] + + K[2][0] * A[5][2] + + K[3][0] * A[5][3] + + K[4][0] * A[5][4] ) * h, ( - K[0, 1] * A[5, 0] - + K[1, 1] * A[5, 1] - + K[2, 1] * A[5, 2] - + K[3, 1] * A[5, 3] - + K[4, 1] * A[5, 4] + K[0][1] * A[5][0] + + K[1][1] * A[5][1] + + K[2][1] * A[5][2] + + K[3][1] * A[5][3] + + K[4][1] * A[5][4] ) * h, ( - K[0, 2] * A[5, 0] - + K[1, 2] * A[5, 1] - + K[2, 2] * A[5, 2] - + K[3, 2] * A[5, 3] - + K[4, 2] * A[5, 4] + K[0][2] * A[5][0] + + K[1][2] * A[5][1] + + K[2][2] * A[5][2] + + K[3][2] * A[5][3] + + K[4][2] * A[5][4] ) * h, ( - K[0, 3] * A[5, 0] - + K[1, 3] * A[5, 1] - + K[2, 3] * A[5, 2] - + K[3, 3] * A[5, 3] - + K[4, 3] * A[5, 4] + K[0][3] * A[5][0] + + K[1][3] * A[5][1] + + K[2][3] * A[5][2] + + K[3][3] * A[5][3] + + K[4][3] * A[5][4] ) * h, ( - K[0, 4] * A[5, 0] - + K[1, 4] * A[5, 1] - + K[2, 4] * A[5, 2] - + K[3, 4] * A[5, 3] - + K[4, 4] * A[5, 4] + K[0][4] * A[5][0] + + K[1][4] * A[5][1] + + K[2][4] * A[5][2] + + K[3][4] * A[5][3] + + K[4][4] * A[5][4] ) * h, ( - K[0, 5] * A[5, 0] - + K[1, 5] * A[5, 1] - + K[2, 5] * A[5, 2] - + K[3, 5] * A[5, 3] - + K[4, 5] * A[5, 4] + K[0][5] * A[5][0] + + K[1][5] * A[5][1] + + K[2][5] * A[5][2] + + K[3][5] * A[5][3] + + K[4][5] * A[5][4] ) * h, ] @@ -259,57 +260,57 @@ def rk_step( dy = np.array( [ ( - K[0, 0] * A[6, 0] - + K[1, 0] * A[6, 1] - + K[2, 0] * A[6, 2] - + K[3, 0] * A[6, 3] - + K[4, 0] * A[6, 4] - + K[5, 0] * A[6, 5] + K[0][0] * A[6][0] + + K[1][0] * A[6][1] + + K[2][0] * A[6][2] + + K[3][0] * A[6][3] + + K[4][0] * A[6][4] + + K[5][0] * A[6][5] ) * h, ( - K[0, 1] * A[6, 0] - + K[1, 1] * A[6, 1] - + K[2, 1] * A[6, 2] - + K[3, 1] * A[6, 3] - + K[4, 1] * A[6, 4] - + K[5, 1] * A[6, 5] + K[0][1] * A[6][0] + + K[1][1] * A[6][1] + + K[2][1] * A[6][2] + + K[3][1] * A[6][3] + + K[4][1] * A[6][4] + + K[5][1] * A[6][5] ) * h, ( - K[0, 2] * A[6, 0] - + K[1, 2] * A[6, 1] - + K[2, 2] * A[6, 2] - + K[3, 2] * A[6, 3] - + K[4, 2] * A[6, 4] - + K[5, 2] * A[6, 5] + K[0][2] * A[6][0] + + K[1][2] * A[6][1] + + K[2][2] * A[6][2] + + K[3][2] * A[6][3] + + K[4][2] * A[6][4] + + K[5][2] * A[6][5] ) * h, ( - K[0, 3] * A[6, 0] - + K[1, 3] * A[6, 1] - + K[2, 3] * A[6, 2] - + K[3, 3] * A[6, 3] - + K[4, 3] * A[6, 4] - + K[5, 3] * A[6, 5] + K[0][3] * A[6][0] + + K[1][3] * A[6][1] + + K[2][3] * A[6][2] + + K[3][3] * A[6][3] + + K[4][3] * A[6][4] + + K[5][3] * A[6][5] ) * h, ( - K[0, 4] * A[6, 0] - + K[1, 4] * A[6, 1] - + K[2, 4] * A[6, 2] - + K[3, 4] * A[6, 3] - + K[4, 4] * A[6, 4] - + K[5, 4] * A[6, 5] + K[0][4] * A[6][0] + + K[1][4] * A[6][1] + + K[2][4] * A[6][2] + + K[3][4] * A[6][3] + + K[4][4] * A[6][4] + + K[5][4] * A[6][5] ) * h, ( - K[0, 5] * A[6, 0] - + K[1, 5] * A[6, 1] - + K[2, 5] * A[6, 2] - + K[3, 5] * A[6, 3] - + K[4, 5] * A[6, 4] - + K[5, 5] * A[6, 5] + K[0][5] * A[6][0] + + K[1][5] * A[6][1] + + K[2][5] * A[6][2] + + K[3][5] * A[6][3] + + K[4][5] * A[6][4] + + K[5][5] * A[6][5] ) * h, ] @@ -320,136 +321,137 @@ def rk_step( dy = np.array( [ ( - K[0, 0] * A[7, 0] - + K[1, 0] * A[7, 1] - + K[2, 0] * A[7, 2] - + K[3, 0] * A[7, 3] - + K[4, 0] * A[7, 4] - + K[5, 0] * A[7, 5] - + K[6, 0] * A[7, 6] + K[0][0] * A[7][0] + + K[1][0] * A[7][1] + + K[2][0] * A[7][2] + + K[3][0] * A[7][3] + + K[4][0] * A[7][4] + + K[5][0] * A[7][5] + + K[6][0] * A[7][6] ) * h, ( - K[0, 1] * A[7, 0] - + K[1, 1] * A[7, 1] - + K[2, 1] * A[7, 2] - + K[3, 1] * A[7, 3] - + K[4, 1] * A[7, 4] - + K[5, 1] * A[7, 5] - + K[6, 1] * A[7, 6] + K[0][1] * A[7][0] + + K[1][1] * A[7][1] + + K[2][1] * A[7][2] + + K[3][1] * A[7][3] + + K[4][1] * A[7][4] + + K[5][1] * A[7][5] + + K[6][1] * A[7][6] ) * h, ( - K[0, 2] * A[7, 0] - + K[1, 2] * A[7, 1] - + K[2, 2] * A[7, 2] - + K[3, 2] * A[7, 3] - + K[4, 2] * A[7, 4] - + K[5, 2] * A[7, 5] - + K[6, 2] * A[7, 6] + K[0][2] * A[7][0] + + K[1][2] * A[7][1] + + K[2][2] * A[7][2] + + K[3][2] * A[7][3] + + K[4][2] * A[7][4] + + K[5][2] * A[7][5] + + K[6][2] * A[7][6] ) * h, ( - K[0, 3] * A[7, 0] - + K[1, 3] * A[7, 1] - + K[2, 3] * A[7, 2] - + K[3, 3] * A[7, 3] - + K[4, 3] * A[7, 4] - + K[5, 3] * A[7, 5] - + K[6, 3] * A[7, 6] + K[0][3] * A[7][0] + + K[1][3] * A[7][1] + + K[2][3] * A[7][2] + + K[3][3] * A[7][3] + + K[4][3] * A[7][4] + + K[5][3] * A[7][5] + + K[6][3] * A[7][6] ) * h, ( - K[0, 4] * A[7, 0] - + K[1, 4] * A[7, 1] - + K[2, 4] * A[7, 2] - + K[3, 4] * A[7, 3] - + K[4, 4] * A[7, 4] - + K[5, 4] * A[7, 5] - + K[6, 4] * A[7, 6] + K[0][4] * A[7][0] + + K[1][4] * A[7][1] + + K[2][4] * A[7][2] + + K[3][4] * A[7][3] + + K[4][4] * A[7][4] + + K[5][4] * A[7][5] + + K[6][4] * A[7][6] ) * h, ( - K[0, 5] * A[7, 0] - + K[1, 5] * A[7, 1] - + K[2, 5] * A[7, 2] - + K[3, 5] * A[7, 3] - + K[4, 5] * A[7, 4] - + K[5, 5] * A[7, 5] - + K[6, 5] * A[7, 6] + K[0][5] * A[7][0] + + K[1][5] * A[7][1] + + K[2][5] * A[7][2] + + K[3][5] * A[7][3] + + K[4][5] * A[7][4] + + K[5][5] * A[7][5] + + K[6][5] * A[7][6] ) * h, ] ) + K[7] = fun(t + C[7] * h, y + dy, argk) # dy = np.dot(K[:8].T, A[8, :8]) * h dy = np.array( [ ( - K[0, 0] * A[8, 0] - + K[1, 0] * A[8, 1] - + K[2, 0] * A[8, 2] - + K[3, 0] * A[8, 3] - + K[4, 0] * A[8, 4] - + K[5, 0] * A[8, 5] - + K[6, 0] * A[8, 6] - + K[7, 0] * A[8, 7] + K[0][0] * A[8][0] + + K[1][0] * A[8][1] + + K[2][0] * A[8][2] + + K[3][0] * A[8][3] + + K[4][0] * A[8][4] + + K[5][0] * A[8][5] + + K[6][0] * A[8][6] + + K[7][0] * A[8][7] ) * h, ( - K[0, 1] * A[8, 0] - + K[1, 1] * A[8, 1] - + K[2, 1] * A[8, 2] - + K[3, 1] * A[8, 3] - + K[4, 1] * A[8, 4] - + K[5, 1] * A[8, 5] - + K[6, 1] * A[8, 6] - + K[7, 1] * A[8, 7] + K[0][1] * A[8][0] + + K[1][1] * A[8][1] + + K[2][1] * A[8][2] + + K[3][1] * A[8][3] + + K[4][1] * A[8][4] + + K[5][1] * A[8][5] + + K[6][1] * A[8][6] + + K[7][1] * A[8][7] ) * h, ( - K[0, 2] * A[8, 0] - + K[1, 2] * A[8, 1] - + K[2, 2] * A[8, 2] - + K[3, 2] * A[8, 3] - + K[4, 2] * A[8, 4] - + K[5, 2] * A[8, 5] - + K[6, 2] * A[8, 6] - + K[7, 2] * A[8, 7] + K[0][2] * A[8][0] + + K[1][2] * A[8][1] + + K[2][2] * A[8][2] + + K[3][2] * A[8][3] + + K[4][2] * A[8][4] + + K[5][2] * A[8][5] + + K[6][2] * A[8][6] + + K[7][2] * A[8][7] ) * h, ( - K[0, 3] * A[8, 0] - + K[1, 3] * A[8, 1] - + K[2, 3] * A[8, 2] - + K[3, 3] * A[8, 3] - + K[4, 3] * A[8, 4] - + K[5, 3] * A[8, 5] - + K[6, 3] * A[8, 6] - + K[7, 3] * A[8, 7] + K[0][3] * A[8][0] + + K[1][3] * A[8][1] + + K[2][3] * A[8][2] + + K[3][3] * A[8][3] + + K[4][3] * A[8][4] + + K[5][3] * A[8][5] + + K[6][3] * A[8][6] + + K[7][3] * A[8][7] ) * h, ( - K[0, 4] * A[8, 0] - + K[1, 4] * A[8, 1] - + K[2, 4] * A[8, 2] - + K[3, 4] * A[8, 3] - + K[4, 4] * A[8, 4] - + K[5, 4] * A[8, 5] - + K[6, 4] * A[8, 6] - + K[7, 4] * A[8, 7] + K[0][4] * A[8][0] + + K[1][4] * A[8][1] + + K[2][4] * A[8][2] + + K[3][4] * A[8][3] + + K[4][4] * A[8][4] + + K[5][4] * A[8][5] + + K[6][4] * A[8][6] + + K[7][4] * A[8][7] ) * h, ( - K[0, 5] * A[8, 0] - + K[1, 5] * A[8, 1] - + K[2, 5] * A[8, 2] - + K[3, 5] * A[8, 3] - + K[4, 5] * A[8, 4] - + K[5, 5] * A[8, 5] - + K[6, 5] * A[8, 6] - + K[7, 5] * A[8, 7] + K[0][5] * A[8][0] + + K[1][5] * A[8][1] + + K[2][5] * A[8][2] + + K[3][5] * A[8][3] + + K[4][5] * A[8][4] + + K[5][5] * A[8][5] + + K[6][5] * A[8][6] + + K[7][5] * A[8][7] ) * h, ] @@ -460,75 +462,75 @@ def rk_step( dy = np.array( [ ( - K[0, 0] * A[9, 0] - + K[1, 0] * A[9, 1] - + K[2, 0] * A[9, 2] - + K[3, 0] * A[9, 3] - + K[4, 0] * A[9, 4] - + K[5, 0] * A[9, 5] - + K[6, 0] * A[9, 6] - + K[7, 0] * A[9, 7] - + K[8, 0] * A[9, 8] + K[0][0] * A[9][0] + + K[1][0] * A[9][1] + + K[2][0] * A[9][2] + + K[3][0] * A[9][3] + + K[4][0] * A[9][4] + + K[5][0] * A[9][5] + + K[6][0] * A[9][6] + + K[7][0] * A[9][7] + + K[8][0] * A[9][8] ) * h, ( - K[0, 1] * A[9, 0] - + K[1, 1] * A[9, 1] - + K[2, 1] * A[9, 2] - + K[3, 1] * A[9, 3] - + K[4, 1] * A[9, 4] - + K[5, 1] * A[9, 5] - + K[6, 1] * A[9, 6] - + K[7, 1] * A[9, 7] - + K[8, 1] * A[9, 8] + K[0][1] * A[9][0] + + K[1][1] * A[9][1] + + K[2][1] * A[9][2] + + K[3][1] * A[9][3] + + K[4][1] * A[9][4] + + K[5][1] * A[9][5] + + K[6][1] * A[9][6] + + K[7][1] * A[9][7] + + K[8][1] * A[9][8] ) * h, ( - K[0, 2] * A[9, 0] - + K[1, 2] * A[9, 1] - + K[2, 2] * A[9, 2] - + K[3, 2] * A[9, 3] - + K[4, 2] * A[9, 4] - + K[5, 2] * A[9, 5] - + K[6, 2] * A[9, 6] - + K[7, 2] * A[9, 7] - + K[8, 2] * A[9, 8] + K[0][2] * A[9][0] + + K[1][2] * A[9][1] + + K[2][2] * A[9][2] + + K[3][2] * A[9][3] + + K[4][2] * A[9][4] + + K[5][2] * A[9][5] + + K[6][2] * A[9][6] + + K[7][2] * A[9][7] + + K[8][2] * A[9][8] ) * h, ( - K[0, 3] * A[9, 0] - + K[1, 3] * A[9, 1] - + K[2, 3] * A[9, 2] - + K[3, 3] * A[9, 3] - + K[4, 3] * A[9, 4] - + K[5, 3] * A[9, 5] - + K[6, 3] * A[9, 6] - + K[7, 3] * A[9, 7] - + K[8, 3] * A[9, 8] + K[0][3] * A[9][0] + + K[1][3] * A[9][1] + + K[2][3] * A[9][2] + + K[3][3] * A[9][3] + + K[4][3] * A[9][4] + + K[5][3] * A[9][5] + + K[6][3] * A[9][6] + + K[7][3] * A[9][7] + + K[8][3] * A[9][8] ) * h, ( - K[0, 4] * A[9, 0] - + K[1, 4] * A[9, 1] - + K[2, 4] * A[9, 2] - + K[3, 4] * A[9, 3] - + K[4, 4] * A[9, 4] - + K[5, 4] * A[9, 5] - + K[6, 4] * A[9, 6] - + K[7, 4] * A[9, 7] - + K[8, 4] * A[9, 8] + K[0][4] * A[9][0] + + K[1][4] * A[9][1] + + K[2][4] * A[9][2] + + K[3][4] * A[9][3] + + K[4][4] * A[9][4] + + K[5][4] * A[9][5] + + K[6][4] * A[9][6] + + K[7][4] * A[9][7] + + K[8][4] * A[9][8] ) * h, ( - K[0, 5] * A[9, 0] - + K[1, 5] * A[9, 1] - + K[2, 5] * A[9, 2] - + K[3, 5] * A[9, 3] - + K[4, 5] * A[9, 4] - + K[5, 5] * A[9, 5] - + K[6, 5] * A[9, 6] - + K[7, 5] * A[9, 7] - + K[8, 5] * A[9, 8] + K[0][5] * A[9][0] + + K[1][5] * A[9][1] + + K[2][5] * A[9][2] + + K[3][5] * A[9][3] + + K[4][5] * A[9][4] + + K[5][5] * A[9][5] + + K[6][5] * A[9][6] + + K[7][5] * A[9][7] + + K[8][5] * A[9][8] ) * h, ] @@ -539,81 +541,81 @@ def rk_step( dy = np.array( [ ( - K[0, 0] * A[10, 0] - + K[1, 0] * A[10, 1] - + K[2, 0] * A[10, 2] - + K[3, 0] * A[10, 3] - + K[4, 0] * A[10, 4] - + K[5, 0] * A[10, 5] - + K[6, 0] * A[10, 6] - + K[7, 0] * A[10, 7] - + K[8, 0] * A[10, 8] - + K[9, 0] * A[10, 9] + K[0][0] * A[10][0] + + K[1][0] * A[10][1] + + K[2][0] * A[10][2] + + K[3][0] * A[10][3] + + K[4][0] * A[10][4] + + K[5][0] * A[10][5] + + K[6][0] * A[10][6] + + K[7][0] * A[10][7] + + K[8][0] * A[10][8] + + K[9][0] * A[10][9] ) * h, ( - K[0, 1] * A[10, 0] - + K[1, 1] * A[10, 1] - + K[2, 1] * A[10, 2] - + K[3, 1] * A[10, 3] - + K[4, 1] * A[10, 4] - + K[5, 1] * A[10, 5] - + K[6, 1] * A[10, 6] - + K[7, 1] * A[10, 7] - + K[8, 1] * A[10, 8] - + K[9, 1] * A[10, 9] + K[0][1] * A[10][0] + + K[1][1] * A[10][1] + + K[2][1] * A[10][2] + + K[3][1] * A[10][3] + + K[4][1] * A[10][4] + + K[5][1] * A[10][5] + + K[6][1] * A[10][6] + + K[7][1] * A[10][7] + + K[8][1] * A[10][8] + + K[9][1] * A[10][9] ) * h, ( - K[0, 2] * A[10, 0] - + K[1, 2] * A[10, 1] - + K[2, 2] * A[10, 2] - + K[3, 2] * A[10, 3] - + K[4, 2] * A[10, 4] - + K[5, 2] * A[10, 5] - + K[6, 2] * A[10, 6] - + K[7, 2] * A[10, 7] - + K[8, 2] * A[10, 8] - + K[9, 2] * A[10, 9] + K[0][2] * A[10][0] + + K[1][2] * A[10][1] + + K[2][2] * A[10][2] + + K[3][2] * A[10][3] + + K[4][2] * A[10][4] + + K[5][2] * A[10][5] + + K[6][2] * A[10][6] + + K[7][2] * A[10][7] + + K[8][2] * A[10][8] + + K[9][2] * A[10][9] ) * h, ( - K[0, 3] * A[10, 0] - + K[1, 3] * A[10, 1] - + K[2, 3] * A[10, 2] - + K[3, 3] * A[10, 3] - + K[4, 3] * A[10, 4] - + K[5, 3] * A[10, 5] - + K[6, 3] * A[10, 6] - + K[7, 3] * A[10, 7] - + K[8, 3] * A[10, 8] - + K[9, 3] * A[10, 9] + K[0][3] * A[10][0] + + K[1][3] * A[10][1] + + K[2][3] * A[10][2] + + K[3][3] * A[10][3] + + K[4][3] * A[10][4] + + K[5][3] * A[10][5] + + K[6][3] * A[10][6] + + K[7][3] * A[10][7] + + K[8][3] * A[10][8] + + K[9][3] * A[10][9] ) * h, ( - K[0, 4] * A[10, 0] - + K[1, 4] * A[10, 1] - + K[2, 4] * A[10, 2] - + K[3, 4] * A[10, 3] - + K[4, 4] * A[10, 4] - + K[5, 4] * A[10, 5] - + K[6, 4] * A[10, 6] - + K[7, 4] * A[10, 7] - + K[8, 4] * A[10, 8] - + K[9, 4] * A[10, 9] + K[0][4] * A[10][0] + + K[1][4] * A[10][1] + + K[2][4] * A[10][2] + + K[3][4] * A[10][3] + + K[4][4] * A[10][4] + + K[5][4] * A[10][5] + + K[6][4] * A[10][6] + + K[7][4] * A[10][7] + + K[8][4] * A[10][8] + + K[9][4] * A[10][9] ) * h, ( - K[0, 5] * A[10, 0] - + K[1, 5] * A[10, 1] - + K[2, 5] * A[10, 2] - + K[3, 5] * A[10, 3] - + K[4, 5] * A[10, 4] - + K[5, 5] * A[10, 5] - + K[6, 5] * A[10, 6] - + K[7, 5] * A[10, 7] - + K[8, 5] * A[10, 8] - + K[9, 5] * A[10, 9] + K[0][5] * A[10][0] + + K[1][5] * A[10][1] + + K[2][5] * A[10][2] + + K[3][5] * A[10][3] + + K[4][5] * A[10][4] + + K[5][5] * A[10][5] + + K[6][5] * A[10][6] + + K[7][5] * A[10][7] + + K[8][5] * A[10][8] + + K[9][5] * A[10][9] ) * h, ] @@ -624,87 +626,87 @@ def rk_step( dy = np.array( [ ( - K[0, 0] * A[11, 0] - + K[1, 0] * A[11, 1] - + K[2, 0] * A[11, 2] - + K[3, 0] * A[11, 3] - + K[4, 0] * A[11, 4] - + K[5, 0] * A[11, 5] - + K[6, 0] * A[11, 6] - + K[7, 0] * A[11, 7] - + K[8, 0] * A[11, 8] - + K[9, 0] * A[11, 9] - + K[10, 0] * A[11, 10] + K[0][0] * A[11][0] + + K[1][0] * A[11][1] + + K[2][0] * A[11][2] + + K[3][0] * A[11][3] + + K[4][0] * A[11][4] + + K[5][0] * A[11][5] + + K[6][0] * A[11][6] + + K[7][0] * A[11][7] + + K[8][0] * A[11][8] + + K[9][0] * A[11][9] + + K[10][0] * A[11][10] ) * h, ( - K[0, 1] * A[11, 0] - + K[1, 1] * A[11, 1] - + K[2, 1] * A[11, 2] - + K[3, 1] * A[11, 3] - + K[4, 1] * A[11, 4] - + K[5, 1] * A[11, 5] - + K[6, 1] * A[11, 6] - + K[7, 1] * A[11, 7] - + K[8, 1] * A[11, 8] - + K[9, 1] * A[11, 9] - + K[10, 1] * A[11, 10] + K[0][1] * A[11][0] + + K[1][1] * A[11][1] + + K[2][1] * A[11][2] + + K[3][1] * A[11][3] + + K[4][1] * A[11][4] + + K[5][1] * A[11][5] + + K[6][1] * A[11][6] + + K[7][1] * A[11][7] + + K[8][1] * A[11][8] + + K[9][1] * A[11][9] + + K[10][1] * A[11][10] ) * h, ( - K[0, 2] * A[11, 0] - + K[1, 2] * A[11, 1] - + K[2, 2] * A[11, 2] - + K[3, 2] * A[11, 3] - + K[4, 2] * A[11, 4] - + K[5, 2] * A[11, 5] - + K[6, 2] * A[11, 6] - + K[7, 2] * A[11, 7] - + K[8, 2] * A[11, 8] - + K[9, 2] * A[11, 9] - + K[10, 2] * A[11, 10] + K[0][2] * A[11][0] + + K[1][2] * A[11][1] + + K[2][2] * A[11][2] + + K[3][2] * A[11][3] + + K[4][2] * A[11][4] + + K[5][2] * A[11][5] + + K[6][2] * A[11][6] + + K[7][2] * A[11][7] + + K[8][2] * A[11][8] + + K[9][2] * A[11][9] + + K[10][2] * A[11][10] ) * h, ( - K[0, 3] * A[11, 0] - + K[1, 3] * A[11, 1] - + K[2, 3] * A[11, 2] - + K[3, 3] * A[11, 3] - + K[4, 3] * A[11, 4] - + K[5, 3] * A[11, 5] - + K[6, 3] * A[11, 6] - + K[7, 3] * A[11, 7] - + K[8, 3] * A[11, 8] - + K[9, 3] * A[11, 9] - + K[10, 3] * A[11, 10] + K[0][3] * A[11][0] + + K[1][3] * A[11][1] + + K[2][3] * A[11][2] + + K[3][3] * A[11][3] + + K[4][3] * A[11][4] + + K[5][3] * A[11][5] + + K[6][3] * A[11][6] + + K[7][3] * A[11][7] + + K[8][3] * A[11][8] + + K[9][3] * A[11][9] + + K[10][3] * A[11][10] ) * h, ( - K[0, 4] * A[11, 0] - + K[1, 4] * A[11, 1] - + K[2, 4] * A[11, 2] - + K[3, 4] * A[11, 3] - + K[4, 4] * A[11, 4] - + K[5, 4] * A[11, 5] - + K[6, 4] * A[11, 6] - + K[7, 4] * A[11, 7] - + K[8, 4] * A[11, 8] - + K[9, 4] * A[11, 9] - + K[10, 4] * A[11, 10] + K[0][4] * A[11][0] + + K[1][4] * A[11][1] + + K[2][4] * A[11][2] + + K[3][4] * A[11][3] + + K[4][4] * A[11][4] + + K[5][4] * A[11][5] + + K[6][4] * A[11][6] + + K[7][4] * A[11][7] + + K[8][4] * A[11][8] + + K[9][4] * A[11][9] + + K[10][4] * A[11][10] ) * h, ( - K[0, 5] * A[11, 0] - + K[1, 5] * A[11, 1] - + K[2, 5] * A[11, 2] - + K[3, 5] * A[11, 3] - + K[4, 5] * A[11, 4] - + K[5, 5] * A[11, 5] - + K[6, 5] * A[11, 6] - + K[7, 5] * A[11, 7] - + K[8, 5] * A[11, 8] - + K[9, 5] * A[11, 9] - + K[10, 5] * A[11, 10] + K[0][5] * A[11][0] + + K[1][5] * A[11][1] + + K[2][5] * A[11][2] + + K[3][5] * A[11][3] + + K[4][5] * A[11][4] + + K[5][5] * A[11][5] + + K[6][5] * A[11][6] + + K[7][5] * A[11][7] + + K[8][5] * A[11][8] + + K[9][5] * A[11][9] + + K[10][5] * A[11][10] ) * h, ] From 2fb74e59dcc5b8915cf47b4315e16fce570c1a04 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 15 Jan 2024 22:33:31 +0100 Subject: [PATCH 119/346] state parser --- src/hapsira/core/jit.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 2abdc6bac..de7a0d177 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -62,6 +62,9 @@ def _parse_signatures(signature: str, noreturn: bool = False) -> Union[str, List # TODO hope for support of "f[:]" return values in cuda target; 2D/4D vectors? signature = signature.replace("M", "Tuple([V,V,V])") # matrix is a tuple of vectors signature = signature.replace("V", "Tuple([f,f,f])") # vector is a tuple of floats + signature = signature.replace( + "S", "Tuple([f,f,f,f,f,f])" + ) # state, two vectors, is a tuple of floats return [signature.replace("f", dtype) for dtype in PRECISIONS] From 3de598bf247a09f6043411697e600d85235f26d2 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 15 Jan 2024 22:34:53 +0100 Subject: [PATCH 120/346] cowell has its own jit --- src/hapsira/core/propagation/cowell.py | 27 +++++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index a28d33024..9ba2cf175 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -1,11 +1,32 @@ import numpy as np -from hapsira.core.propagation.base import func_twobody - +from ..jit import hjit from ..math.ivp import solve_ivp +from ..propagation.base import func_twobody + + +def cowell_jit(func): + """ + Wrapper for hjit to track funcs for cowell + """ + compiled = hjit("S(f,S,f)")(func) + compiled.cowell = None # for debugging + return compiled + + +def cowell(k, r, v, tofs, rtol=1e-11, events=None, f=func_twobody): + """ + Scalar cowell + f : float + r : ndarray (3,) + v : ndarray (3,) + tofs : ??? + rtol : float ... or also ndarray? + """ + assert hasattr(f, "cowell") + assert isinstance(rtol, float) -def cowell(k, r, v, tofs, rtol=1e-11, *, events=None, f=func_twobody): x, y, z = r vx, vy, vz = v From 06dd537e9655f7c5274ac06bd4b901ec5c044cb9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 16 Jan 2024 10:18:40 +0100 Subject: [PATCH 121/346] jit planetocentric_to_AltAz --- src/hapsira/core/events.py | 4 ++-- src/hapsira/core/util.py | 40 ++++++++++++++++++++------------------ 2 files changed, 23 insertions(+), 21 deletions(-) diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index f86d3c73c..8d7151bca 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -2,7 +2,7 @@ import numpy as np from hapsira.core.elements import coe_rotation_matrix_hf, rv2coe_hf, RV2COE_TOL -from hapsira.core.util import planetocentric_to_AltAz +from hapsira.core.util import planetocentric_to_AltAz_hf from .jit import array_to_V_hf from .math.linalg import norm_hf @@ -129,7 +129,7 @@ def elevation_function(k, u_, phi, theta, R, R_p, H): # Position of satellite with respect to a point on attractor. rho = np.subtract(u_[:3], coords) - rot_matrix = planetocentric_to_AltAz(theta, phi) + rot_matrix = np.array(planetocentric_to_AltAz_hf(theta, phi)) new_rho = rot_matrix @ rho new_rho = new_rho / np.linalg.norm(new_rho) diff --git a/src/hapsira/core/util.py b/src/hapsira/core/util.py index 96a42cb54..c05483c31 100644 --- a/src/hapsira/core/util.py +++ b/src/hapsira/core/util.py @@ -11,7 +11,7 @@ "rotation_matrix_gf", "alinspace", "spherical_to_cartesian", - "planetocentric_to_AltAz", + "planetocentric_to_AltAz_hf", ] @@ -106,8 +106,8 @@ def spherical_to_cartesian(v): return norm_vecs * np.stack((x, y, z), axis=-1) -@jit -def planetocentric_to_AltAz(theta, phi): +@hjit("M(f,f)") +def planetocentric_to_AltAz_hf(theta, phi): r"""Defines transformation matrix to convert from Planetocentric coordinate system to the Altitude-Azimuth system. @@ -127,23 +127,25 @@ def planetocentric_to_AltAz(theta, phi): Returns ------- - t_matrix: numpy.ndarray + t_matrix: tuple[tuple[float,float,float],...] Transformation matrix """ # Transformation matrix for converting planetocentric equatorial coordinates to topocentric horizon system. - t_matrix = np.array( - [ - [-np.sin(theta), np.cos(theta), 0], - [ - -np.sin(phi) * np.cos(theta), - -np.sin(phi) * np.sin(theta), - np.cos(phi), - ], - [ - np.cos(phi) * np.cos(theta), - np.cos(phi) * np.sin(theta), - np.sin(phi), - ], - ] + st = sin(theta) + ct = cos(theta) + sp = sin(phi) + cp = cos(phi) + + return ( + (-st, ct, 0.0), + ( + -sp * ct, + -sp * st, + cp, + ), + ( + cp * ct, + cp * st, + sp, + ), ) - return t_matrix From 83ce58df0ac46882469a83e2f0508f3afdf9cb50 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 16 Jan 2024 10:40:26 +0100 Subject: [PATCH 122/346] deactivate trap --- src/hapsira/core/propagation/cowell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 9ba2cf175..ff0d820cc 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -24,7 +24,7 @@ def cowell(k, r, v, tofs, rtol=1e-11, events=None, f=func_twobody): tofs : ??? rtol : float ... or also ndarray? """ - assert hasattr(f, "cowell") + # assert hasattr(f, "cowell") assert isinstance(rtol, float) x, y, z = r From 86bb3ae41b4dbe299b3edceede5b91d2b734bcf4 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 16 Jan 2024 10:51:13 +0100 Subject: [PATCH 123/346] jit line_of_sight --- src/hapsira/core/events.py | 40 ++++++++++++++++++++++-------- src/hapsira/core/perturbations.py | 4 +-- src/hapsira/twobody/events.py | 4 +-- tests/tests_twobody/test_events.py | 6 ++--- 4 files changed, 36 insertions(+), 18 deletions(-) diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index 8d7151bca..a0a049acb 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -1,11 +1,20 @@ +from math import acos + from numba import njit as jit import numpy as np -from hapsira.core.elements import coe_rotation_matrix_hf, rv2coe_hf, RV2COE_TOL -from hapsira.core.util import planetocentric_to_AltAz_hf +from .elements import coe_rotation_matrix_hf, rv2coe_hf, RV2COE_TOL +from .jit import array_to_V_hf, hjit, gjit +from .math.linalg import norm_hf, matmul_VV_hf +from .util import planetocentric_to_AltAz_hf + -from .jit import array_to_V_hf -from .math.linalg import norm_hf +__all__ = [ + "eclipse_function", + "line_of_sight_hf", + "line_of_sight_gf", + "elevation_function", +] @jit @@ -61,8 +70,8 @@ def eclipse_function(k, u_, r_sec, R_sec, R_primary, umbra=True): return shadow_function -@jit -def line_of_sight(r1, r2, R): +@hjit("f(V,V,f)") +def line_of_sight_hf(r1, r2, R): """Calculates the line of sight condition between two position vectors, r1 and r2. Parameters @@ -81,16 +90,25 @@ def line_of_sight(r1, r2, R): located by r1 and r2, else negative. """ - r1_norm = norm_hf(array_to_V_hf(r1)) - r2_norm = norm_hf(array_to_V_hf(r2)) + r1_norm = norm_hf(r1) + r2_norm = norm_hf(r2) - theta = np.arccos((r1 @ r2) / r1_norm / r2_norm) - theta_1 = np.arccos(R / r1_norm) - theta_2 = np.arccos(R / r2_norm) + theta = acos(matmul_VV_hf(r1, r2) / r1_norm / r2_norm) + theta_1 = acos(R / r1_norm) + theta_2 = acos(R / r2_norm) return (theta_1 + theta_2) - theta +@gjit("void(f[:],f[:],f,f[:])", "(n),(n),()->()") +def line_of_sight_gf(r1, r2, R, delta_theta): + """ + Vectorized line_of_sight + """ + + delta_theta[0] = line_of_sight_hf(array_to_V_hf(r1), array_to_V_hf(r2), R) + + @jit def elevation_function(k, u_, phi, theta, R, R_p, H): """Calculates the elevation angle of an object in orbit with respect to diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index 6006e5103..85942ce43 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from hapsira.core.events import line_of_sight as line_of_sight_fast +from hapsira.core.events import line_of_sight_hf from .jit import array_to_V_hf from .math.linalg import norm_hf @@ -240,5 +240,5 @@ def radiation_pressure(t0, state, k, R, C_R, A_over_m, Wdivc_s, star): r_sat = state[:3] P_s = Wdivc_s / (norm_hf(array_to_V_hf(r_star)) ** 2) - nu = float(line_of_sight_fast(r_sat, r_star, R) > 0) + nu = float(line_of_sight_hf(array_to_V_hf(r_sat), array_to_V_hf(r_star), R) > 0) return -nu * P_s * (C_R * A_over_m) * r_star / norm_hf(array_to_V_hf(r_star)) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index 3da687531..dc931bd5b 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -7,7 +7,7 @@ from hapsira.core.math.linalg import norm_vf from hapsira.core.events import ( eclipse_function as eclipse_function_fast, - line_of_sight as line_of_sight_fast, + line_of_sight_gf, ) from hapsira.core.spheroid_location import ( cartesian_to_ellipsoidal as cartesian_to_ellipsoidal_fast, @@ -281,5 +281,5 @@ def __call__(self, t, u_, k): pos_coord = self._pos_coords.pop(0) if self._pos_coords else self._last_coord # Need to cast `pos_coord` to array since `norm` inside numba only works for arrays, not lists. - delta_angle = line_of_sight_fast(u_[:3], np.array(pos_coord), self._R) + delta_angle = line_of_sight_gf(u_[:3], np.array(pos_coord), self._R) return delta_angle diff --git a/tests/tests_twobody/test_events.py b/tests/tests_twobody/test_events.py index b24e7cb01..76e8dda51 100644 --- a/tests/tests_twobody/test_events.py +++ b/tests/tests_twobody/test_events.py @@ -7,7 +7,7 @@ from hapsira.bodies import Earth from hapsira.constants import H0_earth, rho0_earth -from hapsira.core.events import line_of_sight +from hapsira.core.events import line_of_sight_gf from hapsira.core.perturbations import atmospheric_drag_exponential from hapsira.core.propagation import func_twobody from hapsira.twobody import Orbit @@ -330,8 +330,8 @@ def test_line_of_sight(): r_sun = np.array([122233179, -76150708, 33016374]) << u.km R = Earth.R.to(u.km).value - los = line_of_sight(r1.value, r2.value, R) - los_with_sun = line_of_sight(r1.value, r_sun.value, R) + los = line_of_sight_gf(r1.value, r2.value, R) + los_with_sun = line_of_sight_gf(r1.value, r_sun.value, R) assert los < 0 # No LOS condition. assert los_with_sun >= 0 # LOS condition. From af83602ff4428de0a5b1d2b99b06d1b76fed544f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 16 Jan 2024 10:54:04 +0100 Subject: [PATCH 124/346] cleanup --- src/hapsira/core/perturbations.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index 85942ce43..130277943 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -1,12 +1,21 @@ from numba import njit as jit import numpy as np -from hapsira.core.events import line_of_sight_hf - +from .events import line_of_sight_hf from .jit import array_to_V_hf from .math.linalg import norm_hf +__all__ = [ + "J2_perturbation", + "J3_perturbation", + "atmospheric_drag_exponential", + "atmospheric_drag", + "third_body", + "radiation_pressure", +] + + @jit def J2_perturbation(t0, state, k, J2, R): r"""Calculates J2_perturbation acceleration (km/s2). From fdf616d80391e3dae1adea582765ad60ac4eca45 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 16 Jan 2024 12:00:37 +0100 Subject: [PATCH 125/346] jit J2_perturbation --- ...tural-and-artificial-perturbations.myst.md | 9 +++--- src/hapsira/core/math/linalg.py | 6 ++++ src/hapsira/core/perturbations.py | 26 ++++++++-------- src/hapsira/earth/__init__.py | 17 +++++++---- tests/tests_twobody/test_perturbations.py | 30 ++++++++++++++----- 5 files changed, 58 insertions(+), 30 deletions(-) diff --git a/docs/source/examples/natural-and-artificial-perturbations.myst.md b/docs/source/examples/natural-and-artificial-perturbations.myst.md index 4f595b70b..aafa7dcf1 100644 --- a/docs/source/examples/natural-and-artificial-perturbations.myst.md +++ b/docs/source/examples/natural-and-artificial-perturbations.myst.md @@ -25,10 +25,11 @@ from hapsira.bodies import Earth, Moon from hapsira.constants import rho0_earth, H0_earth from hapsira.core.elements import rv2coe_gf, RV2COE_TOL +from hapsira.core.jit import array_to_V_hf from hapsira.core.perturbations import ( atmospheric_drag_exponential, third_body, - J2_perturbation, + J2_perturbation_hf, ) from hapsira.core.propagation import func_twobody from hapsira.ephem import build_ephem_interpolant @@ -148,8 +149,8 @@ tofs = TimeDelta(np.linspace(0, 48.0 * u.h, num=2000)) def f(t0, state, k): du_kep = func_twobody(t0, state, k) - ax, ay, az = J2_perturbation( - t0, state, k, J2=Earth.J2.value, R=Earth.R.to(u.km).value + ax, ay, az = J2_perturbation_hf( + t0, array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), k, J2=Earth.J2.value, R=Earth.R.to(u.km).value ) du_ad = np.array([0, 0, 0, ax, ay, az]) @@ -299,7 +300,7 @@ from numba import njit as jit # Add @jit for speed! @jit def a_d(t0, state, k, J2, R, C_D, A_over_m, H0, rho0): - return J2_perturbation(t0, state, k, J2, R) + atmospheric_drag_exponential( + return np.array(J2_perturbation_hf(t0, array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), k, J2, R)) + atmospheric_drag_exponential( t0, state, k, R, C_D, A_over_m, H0, rho0 ) ``` diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 985fce29e..f856ba1ed 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -10,6 +10,7 @@ "matmul_VM_hf", "matmul_VV_hf", "mul_Vs_hf", + "mul_VV_hf", "norm_hf", "norm_vf", "sign_hf", @@ -87,6 +88,11 @@ def mul_Vs_hf(v, s): return v[0] * s, v[1] * s, v[2] * s +@hjit("V(V,V)") +def mul_VV_hf(a, b): + return a[0] * b[0], a[1] * b[1], a[2] * b[2] + + @hjit("f(V)") def norm_hf(a): return sqrt(matmul_VV_hf(a, a)) diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index 130277943..cafcb3aeb 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -2,12 +2,12 @@ import numpy as np from .events import line_of_sight_hf -from .jit import array_to_V_hf -from .math.linalg import norm_hf +from .jit import array_to_V_hf, hjit +from .math.linalg import norm_hf, mul_Vs_hf, mul_VV_hf __all__ = [ - "J2_perturbation", + "J2_perturbation_hf", "J3_perturbation", "atmospheric_drag_exponential", "atmospheric_drag", @@ -16,8 +16,8 @@ ] -@jit -def J2_perturbation(t0, state, k, J2, R): +@hjit("V(f,V,V,f,f,f)") +def J2_perturbation_hf(t0, rr, vv, k, J2, R): r"""Calculates J2_perturbation acceleration (km/s2). .. math:: @@ -30,8 +30,10 @@ def J2_perturbation(t0, state, k, J2, R): ---------- t0 : float Current time (s) - state : numpy.ndarray - Six component state vector [x, y, z, vx, vy, vz] (km, km/s). + rr : tuple[float,float,float] + Vector [x, y, z] (km) + vv : tuple[float,float,float] + Vector [vx, vy, vz] (km/s) k : float Standard Gravitational parameter. (km^3/s^2) J2 : float @@ -45,15 +47,13 @@ def J2_perturbation(t0, state, k, J2, R): Howard Curtis, (12.30) """ - r_vec = state[:3] - r = norm_hf(array_to_V_hf(r_vec)) + r = norm_hf(rr) factor = (3.0 / 2.0) * k * J2 * (R**2) / (r**5) - a_x = 5.0 * r_vec[2] ** 2 / r**2 - 1 - a_y = 5.0 * r_vec[2] ** 2 / r**2 - 1 - a_z = 5.0 * r_vec[2] ** 2 / r**2 - 3 - return np.array([a_x, a_y, a_z]) * r_vec * factor + a_base = 5.0 * rr[2] ** 2 / r**2 + a = a_base - 1, a_base - 1, a_base - 3 + return mul_Vs_hf(mul_VV_hf(a, rr), factor) @jit diff --git a/src/hapsira/earth/__init__.py b/src/hapsira/earth/__init__.py index 55865dcff..6fb8e8144 100644 --- a/src/hapsira/earth/__init__.py +++ b/src/hapsira/earth/__init__.py @@ -6,7 +6,8 @@ import numpy as np from hapsira.bodies import Earth -from hapsira.core.perturbations import J2_perturbation +from hapsira.core.jit import array_to_V_hf +from hapsira.core.perturbations import J2_perturbation_hf from hapsira.core.propagation import func_twobody from hapsira.earth.enums import EarthGravity from hapsira.twobody.propagation import CowellPropagator @@ -76,17 +77,21 @@ def propagate(self, tof, atmosphere=None, gravity=None, *args): ad_kwargs: Dict[object, dict] = {} perturbations: Dict[object, dict] = {} - def ad(t0, state, k, perturbations): + def ad(t0, state, k, perturbations): # TODO compile + rr, vv = array_to_V_hf(state[:3]), array_to_V_hf(state[3:]) if perturbations: return np.sum( - [f(t0=t0, state=state, k=k, **p) for f, p in perturbations.items()], + [ + f(t0=t0, rr=rr, vv=vv, k=k, **p) + for f, p in perturbations.items() + ], axis=0, ) else: return np.array([0, 0, 0]) - if gravity is EarthGravity.J2: - perturbations[J2_perturbation] = { + if gravity is EarthGravity.J2: # TODO move into compiled `ad` function + perturbations[J2_perturbation_hf] = { "J2": Earth.J2.value, "R": Earth.R.to_value(u.km), } @@ -96,7 +101,7 @@ def ad(t0, state, k, perturbations): # TODO: This whole function probably needs a refactoring raise NotImplementedError - def f(t0, state, k): + def f(t0, state, k): # TODO compile du_kep = func_twobody(t0, state, k) ax, ay, az = ad(t0, state, k, perturbations) du_ad = np.array([0, 0, 0, ax, ay, az]) diff --git a/tests/tests_twobody/test_perturbations.py b/tests/tests_twobody/test_perturbations.py index 8784fd236..0a989708d 100644 --- a/tests/tests_twobody/test_perturbations.py +++ b/tests/tests_twobody/test_perturbations.py @@ -11,8 +11,9 @@ from hapsira.bodies import Earth, Moon, Sun from hapsira.constants import H0_earth, Wdivc_sun, rho0_earth from hapsira.core.elements import rv2coe_gf, RV2COE_TOL +from hapsira.core.jit import array_to_V_hf from hapsira.core.perturbations import ( - J2_perturbation, + J2_perturbation_hf, J3_perturbation, atmospheric_drag, atmospheric_drag_exponential, @@ -40,8 +41,13 @@ def test_J2_propagation_Earth(): def f(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = J2_perturbation( - t0, u_, k, J2=Earth.J2.value, R=Earth.R.to(u.km).value + ax, ay, az = J2_perturbation_hf( + t0, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), + k, + J2=Earth.J2.value, + R=Earth.R.to(u.km).value, ) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad @@ -123,8 +129,13 @@ def test_J3_propagation_Earth(test_params): def f(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = J2_perturbation( - t0, u_, k, J2=Earth.J2.value, R=Earth.R.to(u.km).value + ax, ay, az = J2_perturbation_hf( + t0, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), + k, + J2=Earth.J2.value, + R=Earth.R.to(u.km).value, ) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad @@ -138,8 +149,13 @@ def f(t0, u_, k): def f_combined(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = J2_perturbation( - t0, u_, k, J2=Earth.J2.value, R=Earth.R.to_value(u.km) + ax, ay, az = J2_perturbation_hf( + t0, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), + k, + J2=Earth.J2.value, + R=Earth.R.to_value(u.km), ) + J3_perturbation(t0, u_, k, J3=Earth.J3.value, R=Earth.R.to_value(u.km)) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad From 9591b673f4907d0abacea97a30aa29ad495bf6c7 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 16 Jan 2024 12:44:37 +0100 Subject: [PATCH 126/346] jit J3_perturbation_hf --- src/hapsira/core/perturbations.py | 25 +++++++++---------- tests/tests_twobody/test_perturbations.py | 29 ++++++++++++++++------- 2 files changed, 33 insertions(+), 21 deletions(-) diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index cafcb3aeb..f9495f9dd 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -8,7 +8,7 @@ __all__ = [ "J2_perturbation_hf", - "J3_perturbation", + "J3_perturbation_hf", "atmospheric_drag_exponential", "atmospheric_drag", "third_body", @@ -56,16 +56,18 @@ def J2_perturbation_hf(t0, rr, vv, k, J2, R): return mul_Vs_hf(mul_VV_hf(a, rr), factor) -@jit -def J3_perturbation(t0, state, k, J3, R): +@hjit("V(f,V,V,f,f,f)") +def J3_perturbation_hf(t0, rr, vv, k, J3, R): r"""Calculates J3_perturbation acceleration (km/s2). Parameters ---------- t0 : float Current time (s) - state : numpy.ndarray - Six component state vector [x, y, z, vx, vy, vz] (km, km/s). + rr : tuple[float,float,float] + Vector [x, y, z] (km) + vv : tuple[float,float,float] + Vector [vx, vy, vz] (km/s) k : float Standard Gravitational parameter. (km^3/s^2) J3 : float @@ -77,19 +79,18 @@ def J3_perturbation(t0, state, k, J3, R): ----- The J3 accounts for the oblateness of the attractor. The formula is given in Howard Curtis, problem 12.8 - This perturbation has not been fully validated, see https://github.com/hapsira/hapsira/pull/398 + This perturbation has not been fully validated, see https://github.com/poliastro/poliastro/pull/398 """ - r_vec = state[:3] - r = norm_hf(array_to_V_hf(r_vec)) + r = norm_hf(rr) factor = (1.0 / 2.0) * k * J3 * (R**3) / (r**5) - cos_phi = r_vec[2] / r + cos_phi = rr[2] / r - a_x = 5.0 * r_vec[0] / r * (7.0 * cos_phi**3 - 3.0 * cos_phi) - a_y = 5.0 * r_vec[1] / r * (7.0 * cos_phi**3 - 3.0 * cos_phi) + a_x = 5.0 * rr[0] / r * (7.0 * cos_phi**3 - 3.0 * cos_phi) + a_y = 5.0 * rr[1] / r * (7.0 * cos_phi**3 - 3.0 * cos_phi) a_z = 3.0 * (35.0 / 3.0 * cos_phi**4 - 10.0 * cos_phi**2 + 1) - return np.array([a_x, a_y, a_z]) * factor + return a_x * factor, a_y * factor, a_z * factor @jit diff --git a/tests/tests_twobody/test_perturbations.py b/tests/tests_twobody/test_perturbations.py index 0a989708d..4a386b16c 100644 --- a/tests/tests_twobody/test_perturbations.py +++ b/tests/tests_twobody/test_perturbations.py @@ -14,7 +14,7 @@ from hapsira.core.jit import array_to_V_hf from hapsira.core.perturbations import ( J2_perturbation_hf, - J3_perturbation, + J3_perturbation_hf, atmospheric_drag, atmospheric_drag_exponential, radiation_pressure, @@ -149,14 +149,25 @@ def f(t0, u_, k): def f_combined(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = J2_perturbation_hf( - t0, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), - k, - J2=Earth.J2.value, - R=Earth.R.to_value(u.km), - ) + J3_perturbation(t0, u_, k, J3=Earth.J3.value, R=Earth.R.to_value(u.km)) + ax, ay, az = np.array( + J2_perturbation_hf( + t0, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), + k, + J2=Earth.J2.value, + R=Earth.R.to_value(u.km), + ) + ) + np.array( + J3_perturbation_hf( + t0, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), + k, + J3=Earth.J3.value, + R=Earth.R.to_value(u.km), + ) + ) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad From 41e3c676bee6688d3ae920b3f051a1ff26aa1076 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 16 Jan 2024 13:43:58 +0100 Subject: [PATCH 127/346] jit atmospheric_drag --- docs/source/examples/detecting-events.myst.md | 7 ++-- ...tural-and-artificial-perturbations.myst.md | 16 ++++---- src/hapsira/core/perturbations.py | 41 ++++++++++--------- tests/tests_twobody/test_events.py | 15 +++++-- tests/tests_twobody/test_perturbations.py | 33 +++++++++++---- 5 files changed, 70 insertions(+), 42 deletions(-) diff --git a/docs/source/examples/detecting-events.myst.md b/docs/source/examples/detecting-events.myst.md index de3a48447..6af490c50 100644 --- a/docs/source/examples/detecting-events.myst.md +++ b/docs/source/examples/detecting-events.myst.md @@ -65,7 +65,8 @@ Let's define some natural perturbation conditions for our orbit so that its alti ```{code-cell} from hapsira.constants import H0_earth, rho0_earth -from hapsira.core.perturbations import atmospheric_drag_exponential +from hapsira.core.jit import array_to_V_hf +from hapsira.core.perturbations import atmospheric_drag_exponential_hf from hapsira.core.propagation import func_twobody R = Earth.R.to_value(u.km) @@ -83,8 +84,8 @@ H0 = H0_earth.to_value(u.km) # km def f(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = atmospheric_drag_exponential( - t0, u_, k, R=R, C_D=C_D, A_over_m=A_over_m, H0=H0, rho0=rho0 + ax, ay, az = atmospheric_drag_exponential_hf( + t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k, R=R, C_D=C_D, A_over_m=A_over_m, H0=H0, rho0=rho0 ) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad diff --git a/docs/source/examples/natural-and-artificial-perturbations.myst.md b/docs/source/examples/natural-and-artificial-perturbations.myst.md index aafa7dcf1..697a8e96c 100644 --- a/docs/source/examples/natural-and-artificial-perturbations.myst.md +++ b/docs/source/examples/natural-and-artificial-perturbations.myst.md @@ -27,7 +27,7 @@ from hapsira.constants import rho0_earth, H0_earth from hapsira.core.elements import rv2coe_gf, RV2COE_TOL from hapsira.core.jit import array_to_V_hf from hapsira.core.perturbations import ( - atmospheric_drag_exponential, + atmospheric_drag_exponential_hf, third_body, J2_perturbation_hf, ) @@ -68,9 +68,9 @@ tofs = TimeDelta(np.linspace(0 * u.h, 100000 * u.s, num=2000)) def f(t0, state, k): du_kep = func_twobody(t0, state, k) - ax, ay, az = atmospheric_drag_exponential( + ax, ay, az = atmospheric_drag_exponential_hf( t0, - state, + array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), k, R=R, C_D=C_D, @@ -300,9 +300,9 @@ from numba import njit as jit # Add @jit for speed! @jit def a_d(t0, state, k, J2, R, C_D, A_over_m, H0, rho0): - return np.array(J2_perturbation_hf(t0, array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), k, J2, R)) + atmospheric_drag_exponential( - t0, state, k, R, C_D, A_over_m, H0, rho0 - ) + return np.array(J2_perturbation_hf(t0, array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), k, J2, R)) + np.array(atmospheric_drag_exponential_hf( + t0, array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), k, R, C_D, A_over_m, H0, rho0 + )) ``` ```{code-cell} ipython3 @@ -337,9 +337,9 @@ rr3, _ = orbit.to_ephem( def f(t0, state, k): du_kep = func_twobody(t0, state, k) - ax, ay, az = atmospheric_drag_exponential( + ax, ay, az = atmospheric_drag_exponential_hf( t0, - state, + array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), k, R=R, C_D=C_D, diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index f9495f9dd..de66e1621 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -1,5 +1,4 @@ -from numba import njit as jit -import numpy as np +from math import exp from .events import line_of_sight_hf from .jit import array_to_V_hf, hjit @@ -9,8 +8,8 @@ __all__ = [ "J2_perturbation_hf", "J3_perturbation_hf", - "atmospheric_drag_exponential", - "atmospheric_drag", + "atmospheric_drag_exponential_hf", + "atmospheric_drag_hf", "third_body", "radiation_pressure", ] @@ -93,8 +92,8 @@ def J3_perturbation_hf(t0, rr, vv, k, J3, R): return a_x * factor, a_y * factor, a_z * factor -@jit -def atmospheric_drag_exponential(t0, state, k, R, C_D, A_over_m, H0, rho0): +@hjit("V(f,V,V,f,f,f,f,f,f)") +def atmospheric_drag_exponential_hf(t0, rr, vv, k, R, C_D, A_over_m, H0, rho0): r"""Calculates atmospheric drag acceleration (km/s2). .. math:: @@ -107,8 +106,10 @@ def atmospheric_drag_exponential(t0, state, k, R, C_D, A_over_m, H0, rho0): ---------- t0 : float Current time (s) - state : numpy.ndarray - Six component state vector [x, y, z, vx, vy, vz] (km, km/s). + rr : tuple[float,float,float] + Vector [x, y, z] (km) + vv : tuple[float,float,float] + Vector [vx, vy, vz] (km/s) k : float Standard Gravitational parameter (km^3/s^2). R : float @@ -130,18 +131,17 @@ def atmospheric_drag_exponential(t0, state, k, R, C_D, A_over_m, H0, rho0): the atmospheric density model is rho(H) = rho0 x exp(-H / H0) """ - H = norm_hf(array_to_V_hf(state[:3])) + H = norm_hf(rr) - v_vec = state[3:] - v = norm_hf(array_to_V_hf(v_vec)) + v = norm_hf(vv) B = C_D * A_over_m - rho = rho0 * np.exp(-(H - R) / H0) + rho = rho0 * exp(-(H - R) / H0) - return -(1.0 / 2.0) * rho * B * v * v_vec + return mul_Vs_hf(vv, -(1.0 / 2.0) * rho * B * v) -@jit -def atmospheric_drag(t0, state, k, C_D, A_over_m, rho): +@hjit("V(f,V,V,f,f,f,f)") +def atmospheric_drag_hf(t0, rr, vv, k, C_D, A_over_m, rho): r"""Calculates atmospheric drag acceleration (km/s2). .. math:: @@ -154,8 +154,10 @@ def atmospheric_drag(t0, state, k, C_D, A_over_m, rho): ---------- t0 : float Current time (s). - state : numpy.ndarray - Six component state vector [x, y, z, vx, vy, vz] (km, km/s). + rr : tuple[float,float,float] + Vector [x, y, z] (km) + vv : tuple[float,float,float] + Vector [vx, vy, vz] (km/s) k : float Standard Gravitational parameter (km^3/s^2) C_D : float @@ -171,11 +173,10 @@ def atmospheric_drag(t0, state, k, C_D, A_over_m, rho): computed by a model from hapsira.earth.atmosphere """ - v_vec = state[3:] - v = norm_hf(array_to_V_hf(v_vec)) + v = norm_hf(vv) B = C_D * A_over_m - return -(1.0 / 2.0) * rho * B * v * v_vec + return mul_Vs_hf(vv, -(1.0 / 2.0) * rho * B * v) def third_body(t0, state, k, k_third, perturbation_body): diff --git a/tests/tests_twobody/test_events.py b/tests/tests_twobody/test_events.py index 76e8dda51..b9fbaf093 100644 --- a/tests/tests_twobody/test_events.py +++ b/tests/tests_twobody/test_events.py @@ -8,7 +8,8 @@ from hapsira.bodies import Earth from hapsira.constants import H0_earth, rho0_earth from hapsira.core.events import line_of_sight_gf -from hapsira.core.perturbations import atmospheric_drag_exponential +from hapsira.core.jit import array_to_V_hf +from hapsira.core.perturbations import atmospheric_drag_exponential_hf from hapsira.core.propagation import func_twobody from hapsira.twobody import Orbit from hapsira.twobody.events import ( @@ -49,8 +50,16 @@ def test_altitude_crossing(): def f(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = atmospheric_drag_exponential( - t0, u_, k, R=R, C_D=C_D, A_over_m=A_over_m, H0=H0, rho0=rho0 + ax, ay, az = atmospheric_drag_exponential_hf( + t0, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), + k, + R=R, + C_D=C_D, + A_over_m=A_over_m, + H0=H0, + rho0=rho0, ) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad diff --git a/tests/tests_twobody/test_perturbations.py b/tests/tests_twobody/test_perturbations.py index 4a386b16c..05940f078 100644 --- a/tests/tests_twobody/test_perturbations.py +++ b/tests/tests_twobody/test_perturbations.py @@ -15,8 +15,8 @@ from hapsira.core.perturbations import ( J2_perturbation_hf, J3_perturbation_hf, - atmospheric_drag, - atmospheric_drag_exponential, + atmospheric_drag_hf, + atmospheric_drag_exponential_hf, radiation_pressure, third_body, # pylint: disable=E1120,E1136 ) @@ -269,8 +269,16 @@ def test_atmospheric_drag_exponential(): def f(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = atmospheric_drag_exponential( - t0, u_, k, R=R, C_D=C_D, A_over_m=A_over_m, H0=H0, rho0=rho0 + ax, ay, az = atmospheric_drag_exponential_hf( + t0, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), + k, + R=R, + C_D=C_D, + A_over_m=A_over_m, + H0=H0, + rho0=rho0, ) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad @@ -311,8 +319,16 @@ def test_atmospheric_demise(): def f(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = atmospheric_drag_exponential( - t0, u_, k, R=R, C_D=C_D, A_over_m=A_over_m, H0=H0, rho0=rho0 + ax, ay, az = atmospheric_drag_exponential_hf( + t0, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), + k, + R=R, + C_D=C_D, + A_over_m=A_over_m, + H0=H0, + rho0=rho0, ) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad @@ -369,9 +385,10 @@ def f(t0, u_, k): H = max(norm(u_[:3]), R) rho = coesa76.density((H - R) * u.km).to_value(u.kg / u.km**3) - ax, ay, az = atmospheric_drag( + ax, ay, az = atmospheric_drag_hf( t0, - u_, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), k, C_D=C_D, A_over_m=A_over_m, From 1f7832a0c91661adc80f26d7883d50ac140db5dc Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 20 Jan 2024 15:16:39 +0100 Subject: [PATCH 128/346] independent implementation of interp1d, prep for jit --- src/hapsira/core/math/interpolate.py | 264 ++++++++++++++++++++++++++- 1 file changed, 262 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/math/interpolate.py b/src/hapsira/core/math/interpolate.py index b4614959a..214480ac3 100644 --- a/src/hapsira/core/math/interpolate.py +++ b/src/hapsira/core/math/interpolate.py @@ -1,5 +1,9 @@ import numpy as np -from scipy.interpolate import interp1d + +from scipy.interpolate import interp1d as _scipy_interp1d + +from numpy import asarray, array +from numpy import searchsorted __all__ = [ "interp1d", @@ -8,9 +12,265 @@ ] +def _check_broadcast_up_to(arr_from, shape_to, name): + """Helper to check that arr_from broadcasts up to shape_to""" + shape_from = arr_from.shape + if len(shape_to) >= len(shape_from): + for t, f in zip(shape_to[::-1], shape_from[::-1]): + if f != 1 and f != t: + break + else: # all checks pass, do the upcasting that we need later + if arr_from.size != 1 and arr_from.shape != shape_to: + arr_from = np.ones(shape_to, arr_from.dtype) * arr_from + return arr_from.ravel() + # at least one check failed + raise ValueError( + f"{name} argument must be able to broadcast up " + f"to shape {shape_to} but had shape {shape_from}" + ) + + +class interp1d: + # __slots__ = ("_y_axis", "_y_extra_shape", "dtype") + + def __init__( + self, + x, + y, + axis=-1, + copy=True, + bounds_error=None, + fill_value=np.nan, + assume_sorted=False, + ): + self._y_axis = axis + self._y_extra_shape = None + self.dtype = None + if y is not None: + self._set_yi(y, xi=x, axis=axis) + + self.bounds_error = bounds_error # used by fill_value setter + self.copy = copy + + x = array(x, copy=self.copy) + y = array(y, copy=self.copy) + + if not assume_sorted: # TODO check with assert? + ind = np.argsort(x, kind="mergesort") + x = x[ind] + y = np.take(y, ind, axis=axis) + + assert x.ndim == 1 + assert y.ndim != 0 + + # Backward compatibility + self.axis = axis % y.ndim + + # Interpolation goes internally along the first axis + self.y = y + self._y = self._reshape_yi(self.y) + self.x = x + del y, x # clean up namespace to prevent misuse; use attributes + + minval = 1 + + # Check if we can delegate to numpy.interp (2x-10x faster). TODO + np_dtypes = (np.dtype(np.float64), np.dtype(int)) + cond = self.x.dtype in np_dtypes and self.y.dtype in np_dtypes + cond = cond and self.y.ndim == 1 + + if cond: + self._call = self.__class__._call_linear_np + else: + self._call = self.__class__._call_linear + + assert len(self.x) >= minval + + self.fill_value = fill_value # calls the setter, can modify bounds_err + + def __call__(self, x): + x, x_shape = self._prepare_x(x) + y = self._evaluate(x) + return self._finish_y(y, x_shape) + + def _prepare_x(self, x): + """Reshape input x array to 1-D""" + x = array(x) + x_shape = x.shape + return x.ravel(), x_shape + + def _reshape_yi(self, yi, check=False): + yi = np.moveaxis(np.asarray(yi), self._y_axis, 0) + if check and yi.shape[1:] != self._y_extra_shape: + ok_shape = "{!r} + (N,) + {!r}".format( + self._y_extra_shape[-self._y_axis :], + self._y_extra_shape[: -self._y_axis], + ) + raise ValueError("Data must be of shape %s" % ok_shape) + return yi.reshape((yi.shape[0], -1)) + + def _finish_y(self, y, x_shape): + """Reshape interpolated y back to an N-D array similar to initial y""" + y = y.reshape(x_shape + self._y_extra_shape) + if self._y_axis != 0 and x_shape != (): + nx = len(x_shape) + ny = len(self._y_extra_shape) + s = ( + list(range(nx, nx + self._y_axis)) + + list(range(nx)) + + list(range(nx + self._y_axis, nx + ny)) + ) + y = y.transpose(s) + return y + + def _set_dtype(self, dtype, union=False): + if np.issubdtype(dtype, np.complexfloating) or np.issubdtype( + self.dtype, np.complexfloating + ): + self.dtype = np.complex128 + else: + if not union or self.dtype != np.complex128: + self.dtype = np.float64 + + def _set_yi(self, yi, xi=None, axis=None): + if axis is None: + axis = self._y_axis + if axis is None: + raise ValueError("no interpolation axis specified") + + yi = np.asarray(yi) + + shape = yi.shape + if shape == (): + shape = (1,) + if xi is not None and shape[axis] != len(xi): + raise ValueError( + "x and y arrays must be equal in length along " "interpolation axis." + ) + + self._y_axis = axis % yi.ndim + self._y_extra_shape = yi.shape[: self._y_axis] + yi.shape[self._y_axis + 1 :] + self.dtype = None + self._set_dtype(yi.dtype) + + @property + def fill_value(self): + """The fill value.""" + # backwards compat: mimic a public attribute + return self._fill_value_orig + + @fill_value.setter + def fill_value(self, fill_value): + broadcast_shape = self.y.shape[: self.axis] + self.y.shape[self.axis + 1 :] + if len(broadcast_shape) == 0: + broadcast_shape = (1,) + # it's either a pair (_below_range, _above_range) or a single value + # for both above and below range + if isinstance(fill_value, tuple) and len(fill_value) == 2: + below_above = [np.asarray(fill_value[0]), np.asarray(fill_value[1])] + names = ("fill_value (below)", "fill_value (above)") + for ii in range(2): + below_above[ii] = _check_broadcast_up_to( + below_above[ii], broadcast_shape, names[ii] + ) + else: + fill_value = np.asarray(fill_value) + below_above = [ + _check_broadcast_up_to(fill_value, broadcast_shape, "fill_value") + ] * 2 + self._fill_value_below, self._fill_value_above = below_above + if self.bounds_error is None: + self.bounds_error = True + # backwards compat: fill_value was a public attr; make it writeable + self._fill_value_orig = fill_value + + def _call_linear_np(self, x_new): + # Note that out-of-bounds values are taken care of in self._evaluate + return np.interp(x_new, self.x, self.y) + + def _call_linear(self, x_new): + # 2. Find where in the original data, the values to interpolate + # would be inserted. + # Note: If x_new[n] == x[m], then m is returned by searchsorted. + x_new_indices = searchsorted(self.x, x_new) + + # 3. Clip x_new_indices so that they are within the range of + # self.x indices and at least 1. Removes mis-interpolation + # of x_new[n] = x[0] + x_new_indices = x_new_indices.clip(1, len(self.x) - 1).astype(int) + + # 4. Calculate the slope of regions that each x_new value falls in. + lo = x_new_indices - 1 + hi = x_new_indices + + x_lo = self.x[lo] + x_hi = self.x[hi] + y_lo = self._y[lo] + y_hi = self._y[hi] + + # Note that the following two expressions rely on the specifics of the + # broadcasting semantics. + slope = (y_hi - y_lo) / (x_hi - x_lo)[:, None] + + # 5. Calculate the actual value for each entry in x_new. + y_new = slope * (x_new - x_lo)[:, None] + y_lo + + return y_new + + def _evaluate(self, x_new): + # 1. Handle values in x_new that are outside of x. Throw error, + # or return a list of mask array indicating the outofbounds values. + # The behavior is set by the bounds_error variable. + x_new = asarray(x_new) + y_new = self._call(self, x_new) + below_bounds, above_bounds = self._check_bounds(x_new) + if len(y_new) > 0: + # Note fill_value must be broadcast up to the proper size + # and flattened to work here + y_new[below_bounds] = self._fill_value_below + y_new[above_bounds] = self._fill_value_above + return y_new + + def _check_bounds(self, x_new): + """Check the inputs for being in the bounds of the interpolated data. + + Parameters + ---------- + x_new : array + + Returns + ------- + out_of_bounds : bool array + The mask on x_new of values that are out of the bounds. + """ + + # If self.bounds_error is True, we raise an error if any x_new values + # fall outside the range of x. Otherwise, we return an array indicating + # which values are outside the boundary region. + below_bounds = x_new < self.x[0] + above_bounds = x_new > self.x[-1] + + if self.bounds_error and below_bounds.any(): + below_bounds_value = x_new[np.argmax(below_bounds)] + raise ValueError( + f"A value ({below_bounds_value}) in x_new is below " + f"the interpolation range's minimum value ({self.x[0]})." + ) + if self.bounds_error and above_bounds.any(): + above_bounds_value = x_new[np.argmax(above_bounds)] + raise ValueError( + f"A value ({above_bounds_value}) in x_new is above " + f"the interpolation range's maximum value ({self.x[-1]})." + ) + + # !! Should we emit a warning if some values are out of bounds? + # !! matlab does not. + return below_bounds, above_bounds + + def spline_interp(y, x, u, *, kind="cubic"): """Interpolates y, sampled at x instants, at u instants using `scipy.interpolate.interp1d`.""" - y_u = interp1d(x, y, kind=kind)(u) + y_u = _scipy_interp1d(x, y, kind=kind)(u) return y_u From 0a557f1906a40296e7c47e4c386407b18ecd3bf9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 20 Jan 2024 16:01:33 +0100 Subject: [PATCH 129/346] cleanup --- src/hapsira/core/math/interpolate.py | 118 ++++++--------------------- 1 file changed, 26 insertions(+), 92 deletions(-) diff --git a/src/hapsira/core/math/interpolate.py b/src/hapsira/core/math/interpolate.py index 214480ac3..74d7f2e5b 100644 --- a/src/hapsira/core/math/interpolate.py +++ b/src/hapsira/core/math/interpolate.py @@ -31,62 +31,47 @@ def _check_broadcast_up_to(arr_from, shape_to, name): class interp1d: - # __slots__ = ("_y_axis", "_y_extra_shape", "dtype") - def __init__( self, x, y, - axis=-1, - copy=True, - bounds_error=None, - fill_value=np.nan, - assume_sorted=False, ): - self._y_axis = axis + self._y_axis = -1 self._y_extra_shape = None - self.dtype = None - if y is not None: - self._set_yi(y, xi=x, axis=axis) - - self.bounds_error = bounds_error # used by fill_value setter - self.copy = copy + self._set_yi(y, xi=x, axis=-1) - x = array(x, copy=self.copy) - y = array(y, copy=self.copy) + x = array(x, copy=True) + y = array(y, copy=True) - if not assume_sorted: # TODO check with assert? - ind = np.argsort(x, kind="mergesort") - x = x[ind] - y = np.take(y, ind, axis=axis) + ind = np.argsort(x, kind="mergesort") + x = x[ind] + y = np.take(y, ind, axis=-1) assert x.ndim == 1 assert y.ndim != 0 # Backward compatibility - self.axis = axis % y.ndim + self.axis = -1 % y.ndim # Interpolation goes internally along the first axis self.y = y self._y = self._reshape_yi(self.y) self.x = x - del y, x # clean up namespace to prevent misuse; use attributes - - minval = 1 - - # Check if we can delegate to numpy.interp (2x-10x faster). TODO - np_dtypes = (np.dtype(np.float64), np.dtype(int)) - cond = self.x.dtype in np_dtypes and self.y.dtype in np_dtypes - cond = cond and self.y.ndim == 1 - if cond: - self._call = self.__class__._call_linear_np - else: - self._call = self.__class__._call_linear + assert len(self.x) >= 1 - assert len(self.x) >= minval - - self.fill_value = fill_value # calls the setter, can modify bounds_err + broadcast_shape = self.y.shape[: self.axis] + self.y.shape[self.axis + 1 :] + if len(broadcast_shape) == 0: + broadcast_shape = (1,) + # it's either a pair (_below_range, _above_range) or a single value + # for both above and below range + fill_value = np.asarray(np.nan) + below_above = [ + _check_broadcast_up_to(fill_value, broadcast_shape, "fill_value") + ] * 2 + self._fill_value_below, self._fill_value_above = below_above + # backwards compat: fill_value was a public attr; make it writeable + self._fill_value_orig = fill_value def __call__(self, x): x, x_shape = self._prepare_x(x) @@ -123,15 +108,6 @@ def _finish_y(self, y, x_shape): y = y.transpose(s) return y - def _set_dtype(self, dtype, union=False): - if np.issubdtype(dtype, np.complexfloating) or np.issubdtype( - self.dtype, np.complexfloating - ): - self.dtype = np.complex128 - else: - if not union or self.dtype != np.complex128: - self.dtype = np.float64 - def _set_yi(self, yi, xi=None, axis=None): if axis is None: axis = self._y_axis @@ -150,45 +126,10 @@ def _set_yi(self, yi, xi=None, axis=None): self._y_axis = axis % yi.ndim self._y_extra_shape = yi.shape[: self._y_axis] + yi.shape[self._y_axis + 1 :] - self.dtype = None - self._set_dtype(yi.dtype) - @property - def fill_value(self): - """The fill value.""" - # backwards compat: mimic a public attribute - return self._fill_value_orig - - @fill_value.setter - def fill_value(self, fill_value): - broadcast_shape = self.y.shape[: self.axis] + self.y.shape[self.axis + 1 :] - if len(broadcast_shape) == 0: - broadcast_shape = (1,) - # it's either a pair (_below_range, _above_range) or a single value - # for both above and below range - if isinstance(fill_value, tuple) and len(fill_value) == 2: - below_above = [np.asarray(fill_value[0]), np.asarray(fill_value[1])] - names = ("fill_value (below)", "fill_value (above)") - for ii in range(2): - below_above[ii] = _check_broadcast_up_to( - below_above[ii], broadcast_shape, names[ii] - ) - else: - fill_value = np.asarray(fill_value) - below_above = [ - _check_broadcast_up_to(fill_value, broadcast_shape, "fill_value") - ] * 2 - self._fill_value_below, self._fill_value_above = below_above - if self.bounds_error is None: - self.bounds_error = True - # backwards compat: fill_value was a public attr; make it writeable - self._fill_value_orig = fill_value - - def _call_linear_np(self, x_new): - # Note that out-of-bounds values are taken care of in self._evaluate - return np.interp(x_new, self.x, self.y) + def _evaluate(self, x_new): + x_new = asarray(x_new) - def _call_linear(self, x_new): # 2. Find where in the original data, the values to interpolate # would be inserted. # Note: If x_new[n] == x[m], then m is returned by searchsorted. @@ -215,20 +156,13 @@ def _call_linear(self, x_new): # 5. Calculate the actual value for each entry in x_new. y_new = slope * (x_new - x_lo)[:, None] + y_lo - return y_new - - def _evaluate(self, x_new): - # 1. Handle values in x_new that are outside of x. Throw error, - # or return a list of mask array indicating the outofbounds values. - # The behavior is set by the bounds_error variable. - x_new = asarray(x_new) - y_new = self._call(self, x_new) below_bounds, above_bounds = self._check_bounds(x_new) if len(y_new) > 0: # Note fill_value must be broadcast up to the proper size # and flattened to work here y_new[below_bounds] = self._fill_value_below y_new[above_bounds] = self._fill_value_above + return y_new def _check_bounds(self, x_new): @@ -250,13 +184,13 @@ def _check_bounds(self, x_new): below_bounds = x_new < self.x[0] above_bounds = x_new > self.x[-1] - if self.bounds_error and below_bounds.any(): + if below_bounds.any(): below_bounds_value = x_new[np.argmax(below_bounds)] raise ValueError( f"A value ({below_bounds_value}) in x_new is below " f"the interpolation range's minimum value ({self.x[0]})." ) - if self.bounds_error and above_bounds.any(): + if above_bounds.any(): above_bounds_value = x_new[np.argmax(above_bounds)] raise ValueError( f"A value ({above_bounds_value}) in x_new is above " From f794b4bacae563a2f22a491edba1e693a8182957 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 20 Jan 2024 16:24:34 +0100 Subject: [PATCH 130/346] asserts --- src/hapsira/core/math/interpolate.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/math/interpolate.py b/src/hapsira/core/math/interpolate.py index 74d7f2e5b..786b06cd8 100644 --- a/src/hapsira/core/math/interpolate.py +++ b/src/hapsira/core/math/interpolate.py @@ -36,6 +36,11 @@ def __init__( x, y, ): + assert x.ndim == 1 + assert y.ndim == 2 + assert y.shape[0] == 3 + assert y.shape[1] == x.shape[0] + self._y_axis = -1 self._y_extra_shape = None self._set_yi(y, xi=x, axis=-1) @@ -43,11 +48,11 @@ def __init__( x = array(x, copy=True) y = array(y, copy=True) - ind = np.argsort(x, kind="mergesort") - x = x[ind] - y = np.take(y, ind, axis=-1) + # ind = np.argsort(x, kind="mergesort") + # x = x[ind] + # y = np.take(y, ind, axis=-1) - assert x.ndim == 1 + # assert x.ndim == 1 assert y.ndim != 0 # Backward compatibility @@ -81,6 +86,7 @@ def __call__(self, x): def _prepare_x(self, x): """Reshape input x array to 1-D""" x = array(x) + assert x.shape == tuple() x_shape = x.shape return x.ravel(), x_shape @@ -106,6 +112,7 @@ def _finish_y(self, y, x_shape): + list(range(nx + self._y_axis, nx + ny)) ) y = y.transpose(s) + assert y.shape == (3,) return y def _set_yi(self, yi, xi=None, axis=None): From 2a184064911b9feaa428b8620b6c1c04118ac75e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 12:45:38 +0100 Subject: [PATCH 131/346] cleanup --- src/hapsira/core/math/interpolate.py | 60 ++++------------------------ 1 file changed, 7 insertions(+), 53 deletions(-) diff --git a/src/hapsira/core/math/interpolate.py b/src/hapsira/core/math/interpolate.py index 786b06cd8..317df3b18 100644 --- a/src/hapsira/core/math/interpolate.py +++ b/src/hapsira/core/math/interpolate.py @@ -41,35 +41,18 @@ def __init__( assert y.shape[0] == 3 assert y.shape[1] == x.shape[0] - self._y_axis = -1 - self._y_extra_shape = None - self._set_yi(y, xi=x, axis=-1) + self._y_axis = 1 + self._y_extra_shape = (3,) - x = array(x, copy=True) - y = array(y, copy=True) + self.y = array(y, copy=True) + self.x = array(x, copy=True) - # ind = np.argsort(x, kind="mergesort") - # x = x[ind] - # y = np.take(y, ind, axis=-1) - - # assert x.ndim == 1 - assert y.ndim != 0 - - # Backward compatibility - self.axis = -1 % y.ndim - - # Interpolation goes internally along the first axis - self.y = y - self._y = self._reshape_yi(self.y) - self.x = x + y = np.moveaxis(np.asarray(y), 1, 0) + self._y = y.reshape((y.shape[0], -1)) assert len(self.x) >= 1 - broadcast_shape = self.y.shape[: self.axis] + self.y.shape[self.axis + 1 :] - if len(broadcast_shape) == 0: - broadcast_shape = (1,) - # it's either a pair (_below_range, _above_range) or a single value - # for both above and below range + broadcast_shape = (3,) fill_value = np.asarray(np.nan) below_above = [ _check_broadcast_up_to(fill_value, broadcast_shape, "fill_value") @@ -90,16 +73,6 @@ def _prepare_x(self, x): x_shape = x.shape return x.ravel(), x_shape - def _reshape_yi(self, yi, check=False): - yi = np.moveaxis(np.asarray(yi), self._y_axis, 0) - if check and yi.shape[1:] != self._y_extra_shape: - ok_shape = "{!r} + (N,) + {!r}".format( - self._y_extra_shape[-self._y_axis :], - self._y_extra_shape[: -self._y_axis], - ) - raise ValueError("Data must be of shape %s" % ok_shape) - return yi.reshape((yi.shape[0], -1)) - def _finish_y(self, y, x_shape): """Reshape interpolated y back to an N-D array similar to initial y""" y = y.reshape(x_shape + self._y_extra_shape) @@ -115,25 +88,6 @@ def _finish_y(self, y, x_shape): assert y.shape == (3,) return y - def _set_yi(self, yi, xi=None, axis=None): - if axis is None: - axis = self._y_axis - if axis is None: - raise ValueError("no interpolation axis specified") - - yi = np.asarray(yi) - - shape = yi.shape - if shape == (): - shape = (1,) - if xi is not None and shape[axis] != len(xi): - raise ValueError( - "x and y arrays must be equal in length along " "interpolation axis." - ) - - self._y_axis = axis % yi.ndim - self._y_extra_shape = yi.shape[: self._y_axis] + yi.shape[self._y_axis + 1 :] - def _evaluate(self, x_new): x_new = asarray(x_new) From 0fea6808e26991106136748cfeb473444f66a247 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 13:56:27 +0100 Subject: [PATCH 132/346] cleanup --- src/hapsira/core/math/interpolate.py | 92 ++++++++-------------------- 1 file changed, 25 insertions(+), 67 deletions(-) diff --git a/src/hapsira/core/math/interpolate.py b/src/hapsira/core/math/interpolate.py index 317df3b18..fc6072c8b 100644 --- a/src/hapsira/core/math/interpolate.py +++ b/src/hapsira/core/math/interpolate.py @@ -1,10 +1,8 @@ import numpy as np +from numba import njit from scipy.interpolate import interp1d as _scipy_interp1d -from numpy import asarray, array -from numpy import searchsorted - __all__ = [ "interp1d", "spline_interp", @@ -12,22 +10,17 @@ ] -def _check_broadcast_up_to(arr_from, shape_to, name): - """Helper to check that arr_from broadcasts up to shape_to""" - shape_from = arr_from.shape - if len(shape_to) >= len(shape_from): - for t, f in zip(shape_to[::-1], shape_from[::-1]): - if f != 1 and f != t: - break - else: # all checks pass, do the upcasting that we need later - if arr_from.size != 1 and arr_from.shape != shape_to: - arr_from = np.ones(shape_to, arr_from.dtype) * arr_from - return arr_from.ravel() - # at least one check failed - raise ValueError( - f"{name} argument must be able to broadcast up " - f"to shape {shape_to} but had shape {shape_from}" - ) +@njit +def bisect_left(a, x): + lo = 0 + hi = len(a) + while lo < hi: + mid = (lo + hi) // 2 + if a[mid] < x: + lo = mid + 1 + else: + hi = mid + return lo class interp1d: @@ -38,63 +31,27 @@ def __init__( ): assert x.ndim == 1 assert y.ndim == 2 + assert x.shape[0] >= 1 # > instead of >= assert y.shape[0] == 3 assert y.shape[1] == x.shape[0] - self._y_axis = 1 - self._y_extra_shape = (3,) - - self.y = array(y, copy=True) - self.x = array(x, copy=True) - - y = np.moveaxis(np.asarray(y), 1, 0) - self._y = y.reshape((y.shape[0], -1)) + self.y = y.T + self.x = np.array(x, copy=True) - assert len(self.x) >= 1 - - broadcast_shape = (3,) - fill_value = np.asarray(np.nan) - below_above = [ - _check_broadcast_up_to(fill_value, broadcast_shape, "fill_value") - ] * 2 - self._fill_value_below, self._fill_value_above = below_above - # backwards compat: fill_value was a public attr; make it writeable - self._fill_value_orig = fill_value + self._fill_value_below = np.array([np.nan]) + self._fill_value_above = np.array([np.nan]) def __call__(self, x): - x, x_shape = self._prepare_x(x) - y = self._evaluate(x) - return self._finish_y(y, x_shape) - - def _prepare_x(self, x): - """Reshape input x array to 1-D""" - x = array(x) - assert x.shape == tuple() - x_shape = x.shape - return x.ravel(), x_shape - - def _finish_y(self, y, x_shape): - """Reshape interpolated y back to an N-D array similar to initial y""" - y = y.reshape(x_shape + self._y_extra_shape) - if self._y_axis != 0 and x_shape != (): - nx = len(x_shape) - ny = len(self._y_extra_shape) - s = ( - list(range(nx, nx + self._y_axis)) - + list(range(nx)) - + list(range(nx + self._y_axis, nx + ny)) - ) - y = y.transpose(s) - assert y.shape == (3,) - return y + "x is scalar" - def _evaluate(self, x_new): - x_new = asarray(x_new) + x_new = np.array([x]) # 2. Find where in the original data, the values to interpolate # would be inserted. # Note: If x_new[n] == x[m], then m is returned by searchsorted. - x_new_indices = searchsorted(self.x, x_new) + x_new_indices = np.array( + [bisect_left(self.x, x_new[0])] + ) # np.searchsorted(self.x, x_new) # 3. Clip x_new_indices so that they are within the range of # self.x indices and at least 1. Removes mis-interpolation @@ -107,8 +64,8 @@ def _evaluate(self, x_new): x_lo = self.x[lo] x_hi = self.x[hi] - y_lo = self._y[lo] - y_hi = self._y[hi] + y_lo = self.y[lo] + y_hi = self.y[hi] # Note that the following two expressions rely on the specifics of the # broadcasting semantics. @@ -124,6 +81,7 @@ def _evaluate(self, x_new): y_new[below_bounds] = self._fill_value_below y_new[above_bounds] = self._fill_value_above + y_new = y_new.reshape((3,)) return y_new def _check_bounds(self, x_new): From 70a6eeb31a09bf5a359e442730318da62383dd41 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 16:18:25 +0100 Subject: [PATCH 133/346] jit interp_hf --- src/hapsira/core/math/interpolate.py | 165 ++++++++-------------- src/hapsira/core/perturbations.py | 6 +- src/hapsira/ephem.py | 4 +- tests/tests_twobody/test_perturbations.py | 20 ++- 4 files changed, 70 insertions(+), 125 deletions(-) diff --git a/src/hapsira/core/math/interpolate.py b/src/hapsira/core/math/interpolate.py index fc6072c8b..a70e654f5 100644 --- a/src/hapsira/core/math/interpolate.py +++ b/src/hapsira/core/math/interpolate.py @@ -1,124 +1,69 @@ -import numpy as np -from numba import njit +from typing import Callable +import numpy as np from scipy.interpolate import interp1d as _scipy_interp1d +from ..jit import hjit +from .linalg import add_VV_hf, div_Vs_hf, mul_Vs_hf, sub_VV_hf + __all__ = [ - "interp1d", + "interp_hb", "spline_interp", "sinc_interp", ] -@njit -def bisect_left(a, x): - lo = 0 - hi = len(a) - while lo < hi: - mid = (lo + hi) // 2 - if a[mid] < x: - lo = mid + 1 - else: - hi = mid - return lo - - -class interp1d: - def __init__( - self, - x, - y, - ): - assert x.ndim == 1 - assert y.ndim == 2 - assert x.shape[0] >= 1 # > instead of >= - assert y.shape[0] == 3 - assert y.shape[1] == x.shape[0] - - self.y = y.T - self.x = np.array(x, copy=True) - - self._fill_value_below = np.array([np.nan]) - self._fill_value_above = np.array([np.nan]) - - def __call__(self, x): - "x is scalar" - - x_new = np.array([x]) - - # 2. Find where in the original data, the values to interpolate - # would be inserted. - # Note: If x_new[n] == x[m], then m is returned by searchsorted. - x_new_indices = np.array( - [bisect_left(self.x, x_new[0])] - ) # np.searchsorted(self.x, x_new) - - # 3. Clip x_new_indices so that they are within the range of - # self.x indices and at least 1. Removes mis-interpolation - # of x_new[n] = x[0] - x_new_indices = x_new_indices.clip(1, len(self.x) - 1).astype(int) - - # 4. Calculate the slope of regions that each x_new value falls in. - lo = x_new_indices - 1 - hi = x_new_indices - - x_lo = self.x[lo] - x_hi = self.x[hi] - y_lo = self.y[lo] - y_hi = self.y[hi] - - # Note that the following two expressions rely on the specifics of the - # broadcasting semantics. - slope = (y_hi - y_lo) / (x_hi - x_lo)[:, None] - - # 5. Calculate the actual value for each entry in x_new. - y_new = slope * (x_new - x_lo)[:, None] + y_lo - - below_bounds, above_bounds = self._check_bounds(x_new) - if len(y_new) > 0: - # Note fill_value must be broadcast up to the proper size - # and flattened to work here - y_new[below_bounds] = self._fill_value_below - y_new[above_bounds] = self._fill_value_above - - y_new = y_new.reshape((3,)) +def interp_hb(x: np.ndarray, y: np.ndarray) -> Callable: + """ + Build compiled 1d-interpolator for 1D vectors + """ + + assert x.ndim == 1 + assert y.ndim == 2 + assert x.shape[0] >= 1 # > instead of >= + assert y.shape[0] == 3 + assert y.shape[1] == x.shape[0] + + y = tuple(tuple(record) for record in y.T) + x = tuple(x) + x_len = len(x) + + @hjit("V(f)") + def interp_hf(x_new): + assert x_new >= x[0] + assert x_new <= x[-1] + + # bisect left + x_new_index = 0 + hi = x_len + while x_new_index < hi: + mid = (x_new_index + hi) // 2 + if x[mid] < x_new: + x_new_index = mid + 1 + else: + hi = mid + + # clip + if x_new_index > x_len: + x_new_index = x_len + if x_new_index < 1: + x_new_index = 1 + + # slope + lo = x_new_index - 1 + hi = x_new_index + x_lo = x[lo] + x_hi = x[hi] + y_lo = y[lo] # tuple + y_hi = y[hi] # tuple + slope = div_Vs_hf(sub_VV_hf(y_hi, y_lo), x_hi - x_lo) # tuple + + # new value + y_new = add_VV_hf(mul_Vs_hf(slope, x_new - x_lo), y_lo) # tuple + return y_new - def _check_bounds(self, x_new): - """Check the inputs for being in the bounds of the interpolated data. - - Parameters - ---------- - x_new : array - - Returns - ------- - out_of_bounds : bool array - The mask on x_new of values that are out of the bounds. - """ - - # If self.bounds_error is True, we raise an error if any x_new values - # fall outside the range of x. Otherwise, we return an array indicating - # which values are outside the boundary region. - below_bounds = x_new < self.x[0] - above_bounds = x_new > self.x[-1] - - if below_bounds.any(): - below_bounds_value = x_new[np.argmax(below_bounds)] - raise ValueError( - f"A value ({below_bounds_value}) in x_new is below " - f"the interpolation range's minimum value ({self.x[0]})." - ) - if above_bounds.any(): - above_bounds_value = x_new[np.argmax(above_bounds)] - raise ValueError( - f"A value ({above_bounds_value}) in x_new is above " - f"the interpolation range's maximum value ({self.x[-1]})." - ) - - # !! Should we emit a warning if some values are out of bounds? - # !! matlab does not. - return below_bounds, above_bounds + return interp_hf def spline_interp(y, x, u, *, kind="cubic"): diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index de66e1621..7bfbe4723 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -1,5 +1,7 @@ from math import exp +import numpy as np + from .events import line_of_sight_hf from .jit import array_to_V_hf, hjit from .math.linalg import norm_hf, mul_Vs_hf, mul_VV_hf @@ -206,7 +208,7 @@ def third_body(t0, state, k, k_third, perturbation_body): the gravity from the Moon acting on a small satellite. """ - body_r = perturbation_body(t0) + body_r = np.array(perturbation_body(t0)) delta_r = body_r - state[:3] return ( k_third * delta_r / norm_hf(array_to_V_hf(delta_r)) ** 3 @@ -247,7 +249,7 @@ def radiation_pressure(t0, state, k, R, C_R, A_over_m, Wdivc_s, star): Howard Curtis, section 12.9 """ - r_star = star(t0) + r_star = np.array(star(t0)) r_sat = state[:3] P_s = Wdivc_s / (norm_hf(array_to_V_hf(r_star)) ** 2) diff --git a/src/hapsira/ephem.py b/src/hapsira/ephem.py index 853583028..38705e3d5 100644 --- a/src/hapsira/ephem.py +++ b/src/hapsira/ephem.py @@ -11,7 +11,7 @@ from astroquery.jplhorizons import Horizons from hapsira.bodies import Earth -from hapsira.core.math.interpolate import interp1d, sinc_interp, spline_interp +from hapsira.core.math.interpolate import interp_hb, sinc_interp, spline_interp from hapsira.frames import Planes from hapsira.frames.util import get_frame from hapsira.twobody.sampling import EpochsArray @@ -43,7 +43,7 @@ def build_ephem_interpolant(body, epochs, attractor=Earth): """ ephem = Ephem.from_body(body, epochs, attractor=attractor) - interpolant = interp1d( + interpolant = interp_hb( (epochs - epochs[0]).to_value(u.s), ephem._coordinates.xyz.to_value(u.km), ) diff --git a/tests/tests_twobody/test_perturbations.py b/tests/tests_twobody/test_perturbations.py index 05940f078..9393cf8e6 100644 --- a/tests/tests_twobody/test_perturbations.py +++ b/tests/tests_twobody/test_perturbations.py @@ -1,5 +1,3 @@ -import functools - from astropy import units as u from astropy.coordinates import Angle from astropy.tests.helper import assert_quantity_allclose @@ -11,7 +9,8 @@ from hapsira.bodies import Earth, Moon, Sun from hapsira.constants import H0_earth, Wdivc_sun, rho0_earth from hapsira.core.elements import rv2coe_gf, RV2COE_TOL -from hapsira.core.jit import array_to_V_hf +from hapsira.core.jit import array_to_V_hf, hjit +from hapsira.core.math.linalg import mul_Vs_hf, norm_hf from hapsira.core.perturbations import ( J2_perturbation_hf, J3_perturbation_hf, @@ -665,12 +664,7 @@ def sun_r(): tof = 600 * u.day epoch = Time(j_date, format="jd", scale="tdb") ephem_epochs = time_range(epoch, num_values=164, end=epoch + tof) - return build_ephem_interpolant(Sun, ephem_epochs) - - -def normalize_to_Curtis(t0, sun_r): - r = sun_r(t0) - return 149600000 * r / norm(r) + return build_ephem_interpolant(Sun, ephem_epochs) # returns hf @pytest.mark.slow @@ -702,8 +696,12 @@ def test_solar_pressure(t_days, deltas_expected, sun_r): nu=343.4268 * u.deg, epoch=epoch, ) + # In Curtis, the mean distance to Sun is used. In order to validate against it, we have to do the same thing - sun_normalized = functools.partial(normalize_to_Curtis, sun_r=sun_r) + @hjit("V(f)") + def sun_normalized_hf(t0): + r = sun_r(t0) # sun_r is hf, returns V + return mul_Vs_hf(r, 149600000 / norm_hf(r)) def f(t0, u_, k): du_kep = func_twobody(t0, u_, k) @@ -715,7 +713,7 @@ def f(t0, u_, k): C_R=2.0, A_over_m=2e-4 / 100, Wdivc_s=Wdivc_sun.value, - star=sun_normalized, + star=sun_normalized_hf, ) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad From a05f5be528544fc897ee2b0147bafe8f1263e65d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 16:42:56 +0100 Subject: [PATCH 134/346] jit third_body --- ...tural-and-artificial-perturbations.myst.md | 7 +++--- src/hapsira/core/jit.py | 6 +++-- src/hapsira/core/perturbations.py | 23 +++++++++++-------- tests/tests_twobody/test_perturbations.py | 11 +++++---- 4 files changed, 27 insertions(+), 20 deletions(-) diff --git a/docs/source/examples/natural-and-artificial-perturbations.myst.md b/docs/source/examples/natural-and-artificial-perturbations.myst.md index 697a8e96c..570cea410 100644 --- a/docs/source/examples/natural-and-artificial-perturbations.myst.md +++ b/docs/source/examples/natural-and-artificial-perturbations.myst.md @@ -28,7 +28,7 @@ from hapsira.core.elements import rv2coe_gf, RV2COE_TOL from hapsira.core.jit import array_to_V_hf from hapsira.core.perturbations import ( atmospheric_drag_exponential_hf, - third_body, + third_body_hf, J2_perturbation_hf, ) from hapsira.core.propagation import func_twobody @@ -210,9 +210,10 @@ tofs = TimeDelta(np.linspace(0, 60 * u.day, num=1000)) def f(t0, state, k): du_kep = func_twobody(t0, state, k) - ax, ay, az = third_body( + ax, ay, az = third_body_hf( t0, - state, + array_to_V_hf(state[:3]), + array_to_V_hf(state[3:]), k, k_third=400 * Moon.k.to(u.km**3 / u.s**2).value, perturbation_body=body_r, diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index de7a0d177..1232475e2 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -59,12 +59,14 @@ def _parse_signatures(signature: str, noreturn: bool = False) -> Union[str, List ) return signature - # TODO hope for support of "f[:]" return values in cuda target; 2D/4D vectors? signature = signature.replace("M", "Tuple([V,V,V])") # matrix is a tuple of vectors signature = signature.replace("V", "Tuple([f,f,f])") # vector is a tuple of floats signature = signature.replace( "S", "Tuple([f,f,f,f,f,f])" - ) # state, two vectors, is a tuple of floats + ) # state, two vectors, is a tuple of floats TODO remove, use vectors instead? + signature = signature.replace( + "F", "FunctionType" + ) # TODO does not work for CUDA yet return [signature.replace("f", dtype) for dtype in PRECISIONS] diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index 7bfbe4723..edbe29128 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -4,7 +4,7 @@ from .events import line_of_sight_hf from .jit import array_to_V_hf, hjit -from .math.linalg import norm_hf, mul_Vs_hf, mul_VV_hf +from .math.linalg import norm_hf, mul_Vs_hf, mul_VV_hf, sub_VV_hf __all__ = [ @@ -12,7 +12,7 @@ "J3_perturbation_hf", "atmospheric_drag_exponential_hf", "atmospheric_drag_hf", - "third_body", + "third_body_hf", "radiation_pressure", ] @@ -181,7 +181,8 @@ def atmospheric_drag_hf(t0, rr, vv, k, C_D, A_over_m, rho): return mul_Vs_hf(vv, -(1.0 / 2.0) * rho * B * v) -def third_body(t0, state, k, k_third, perturbation_body): +@hjit("V(f,V,V,f,f,F(V(f)))") +def third_body_hf(t0, rr, vv, k, k_third, perturbation_body): r"""Calculate third body acceleration (km/s2). .. math:: @@ -192,8 +193,10 @@ def third_body(t0, state, k, k_third, perturbation_body): ---------- t0 : float Current time (s). - state : numpy.ndarray - Six component state vector [x, y, z, vx, vy, vz] (km, km/s). + rr : tuple[float,float,float] + Vector [x, y, z] (km) + vv : tuple[float,float,float] + Vector [vx, vy, vz] (km/s) k : float Standard Gravitational parameter of the attractor (km^3/s^2). k_third : float @@ -208,11 +211,11 @@ def third_body(t0, state, k, k_third, perturbation_body): the gravity from the Moon acting on a small satellite. """ - body_r = np.array(perturbation_body(t0)) - delta_r = body_r - state[:3] - return ( - k_third * delta_r / norm_hf(array_to_V_hf(delta_r)) ** 3 - - k_third * body_r / norm_hf(array_to_V_hf(body_r)) ** 3 + body_r = perturbation_body(t0) + delta_r = sub_VV_hf(body_r, rr) + return sub_VV_hf( + mul_Vs_hf(delta_r, k_third / norm_hf(delta_r) ** 3), + mul_Vs_hf(body_r, k_third / norm_hf(body_r) ** 3), ) diff --git a/tests/tests_twobody/test_perturbations.py b/tests/tests_twobody/test_perturbations.py index 9393cf8e6..dd9448255 100644 --- a/tests/tests_twobody/test_perturbations.py +++ b/tests/tests_twobody/test_perturbations.py @@ -17,7 +17,7 @@ atmospheric_drag_hf, atmospheric_drag_exponential_hf, radiation_pressure, - third_body, # pylint: disable=E1120,E1136 + third_body_hf, # pylint: disable=E1120,E1136 ) from hapsira.core.propagation import func_twobody from hapsira.earth.atmosphere import COESA76 @@ -611,12 +611,13 @@ def test_3rd_body_Curtis(test_params): def f(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = third_body( + ax, ay, az = third_body_hf( t0, - u_, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), k, - k_third=body.k.to_value(u.km**3 / u.s**2), - perturbation_body=body_r, + body.k.to_value(u.km**3 / u.s**2), # k_third + body_r, # perturbation_body ) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad From 691a449476495c1e56d21eb358c26c12d341c91c Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 17:07:37 +0100 Subject: [PATCH 135/346] jit radiation_pressure --- src/hapsira/core/perturbations.py | 29 +++++++++++++---------- tests/tests_twobody/test_perturbations.py | 11 +++++---- 2 files changed, 22 insertions(+), 18 deletions(-) diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index edbe29128..ae9ecbbee 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -1,9 +1,7 @@ from math import exp -import numpy as np - from .events import line_of_sight_hf -from .jit import array_to_V_hf, hjit +from .jit import hjit from .math.linalg import norm_hf, mul_Vs_hf, mul_VV_hf, sub_VV_hf @@ -13,7 +11,7 @@ "atmospheric_drag_exponential_hf", "atmospheric_drag_hf", "third_body_hf", - "radiation_pressure", + "radiation_pressure_hf", ] @@ -219,7 +217,8 @@ def third_body_hf(t0, rr, vv, k, k_third, perturbation_body): ) -def radiation_pressure(t0, state, k, R, C_R, A_over_m, Wdivc_s, star): +@hjit("V(f,V,V,f,f,f,f,f,F(V(f)))") +def radiation_pressure_hf(t0, rr, vv, k, R, C_R, A_over_m, Wdivc_s, star): r"""Calculates radiation pressure acceleration (km/s2). .. math:: @@ -230,8 +229,10 @@ def radiation_pressure(t0, state, k, R, C_R, A_over_m, Wdivc_s, star): ---------- t0 : float Current time (s). - state : numpy.ndarray - Six component state vector [x, y, z, vx, vy, vz] (km, km/s). + rr : tuple[float,float,float] + Vector [x, y, z] (km) + vv : tuple[float,float,float] + Vector [vx, vy, vz] (km/s) k : float Standard Gravitational parameter (km^3/s^2). R : float @@ -252,9 +253,11 @@ def radiation_pressure(t0, state, k, R, C_R, A_over_m, Wdivc_s, star): Howard Curtis, section 12.9 """ - r_star = np.array(star(t0)) - r_sat = state[:3] - P_s = Wdivc_s / (norm_hf(array_to_V_hf(r_star)) ** 2) - - nu = float(line_of_sight_hf(array_to_V_hf(r_sat), array_to_V_hf(r_star), R) > 0) - return -nu * P_s * (C_R * A_over_m) * r_star / norm_hf(array_to_V_hf(r_star)) + r_star = star(t0) + P_s = Wdivc_s / (norm_hf(r_star) ** 2) + + if line_of_sight_hf(rr, r_star, R) > 0: + nu = 1.0 + else: + nu = 0.0 + return mul_Vs_hf(r_star, -nu * P_s * (C_R * A_over_m) / norm_hf(r_star)) diff --git a/tests/tests_twobody/test_perturbations.py b/tests/tests_twobody/test_perturbations.py index dd9448255..09c6bdd94 100644 --- a/tests/tests_twobody/test_perturbations.py +++ b/tests/tests_twobody/test_perturbations.py @@ -11,13 +11,13 @@ from hapsira.core.elements import rv2coe_gf, RV2COE_TOL from hapsira.core.jit import array_to_V_hf, hjit from hapsira.core.math.linalg import mul_Vs_hf, norm_hf -from hapsira.core.perturbations import ( +from hapsira.core.perturbations import ( # pylint: disable=E1120,E1136 J2_perturbation_hf, J3_perturbation_hf, atmospheric_drag_hf, atmospheric_drag_exponential_hf, - radiation_pressure, - third_body_hf, # pylint: disable=E1120,E1136 + radiation_pressure_hf, + third_body_hf, ) from hapsira.core.propagation import func_twobody from hapsira.earth.atmosphere import COESA76 @@ -706,9 +706,10 @@ def sun_normalized_hf(t0): def f(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = radiation_pressure( + ax, ay, az = radiation_pressure_hf( t0, - u_, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), k, R=Earth.R.to(u.km).value, C_R=2.0, From d6ac5cd45ede0fcee21afb264c2fc3adb698bffc Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 18:09:09 +0100 Subject: [PATCH 136/346] jit change_a_inc --- src/hapsira/core/thrust/__init__.py | 12 ++- src/hapsira/core/thrust/change_a_inc.py | 98 +++++++++++----------- src/hapsira/twobody/thrust/__init__.py | 10 +-- src/hapsira/twobody/thrust/change_a_inc.py | 12 +-- tests/tests_twobody/test_thrust.py | 11 +-- 5 files changed, 74 insertions(+), 69 deletions(-) diff --git a/src/hapsira/core/thrust/__init__.py b/src/hapsira/core/thrust/__init__.py index 873a4c3f7..11a0ab926 100644 --- a/src/hapsira/core/thrust/__init__.py +++ b/src/hapsira/core/thrust/__init__.py @@ -1,5 +1,9 @@ -from hapsira.core.thrust.change_a_inc import change_a_inc -from hapsira.core.thrust.change_argp import change_argp -from hapsira.core.thrust.change_ecc_inc import change_ecc_inc +# from .change_a_inc import change_a_inc_hb +from .change_argp import change_argp +from .change_ecc_inc import change_ecc_inc -__all__ = ["change_a_inc", "change_argp", "change_ecc_inc"] +__all__ = [ + # "change_a_inc_hb", + "change_argp", + "change_ecc_inc", +] diff --git a/src/hapsira/core/thrust/change_a_inc.py b/src/hapsira/core/thrust/change_a_inc.py index cbd166ef2..202615c39 100644 --- a/src/hapsira/core/thrust/change_a_inc.py +++ b/src/hapsira/core/thrust/change_a_inc.py @@ -1,61 +1,61 @@ -from numba import njit as jit -import numpy as np -from numpy import cross +from math import atan2, cos, pi, sin, tan -from hapsira.core.elements import circular_velocity_hf +from ..jit import hjit +from ..elements import circular_velocity_hf +from ..math.linalg import add_VV_hf, cross_VV_hf, div_Vs_hf, mul_Vs_hf, norm_hf, sign_hf -from ..jit import array_to_V_hf -from ..math.linalg import norm_hf - -@jit -def extra_quantities(k, a_0, a_f, inc_0, inc_f, f): - """Extra quantities given by the Edelbaum (a, i) model.""" - V_0, V_f, beta_0_ = compute_parameters(k, a_0, a_f, inc_0, inc_f) - delta_V_ = delta_V(V_0, V_f, beta_0_, inc_0, inc_f) - t_f_ = delta_V_ / f - - return delta_V_, t_f_ - - -@jit -def beta(t, V_0, f, beta_0): - """Compute yaw angle (β) as a function of time and the problem parameters.""" - return np.arctan2(V_0 * np.sin(beta_0), V_0 * np.cos(beta_0) - f * t) +__all__ = [ + "change_a_inc_hb", +] -@jit -def beta_0(V_0, V_f, inc_0, inc_f): +@hjit("f(f,f,f,f)") +def _beta_0_hf(V_0, V_f, inc_0, inc_f): """Compute initial yaw angle (β) as a function of the problem parameters.""" delta_i_f = abs(inc_f - inc_0) - return np.arctan2( - np.sin(np.pi / 2 * delta_i_f), - V_0 / V_f - np.cos(np.pi / 2 * delta_i_f), + return atan2( + sin(pi / 2 * delta_i_f), + V_0 / V_f - cos(pi / 2 * delta_i_f), ) -@jit -def compute_parameters(k, a_0, a_f, inc_0, inc_f): +@hjit("Tuple([f,f,f])(f,f,f,f,f)") +def _compute_parameters_hf(k, a_0, a_f, inc_0, inc_f): """Compute parameters of the model.""" V_0 = circular_velocity_hf(k, a_0) V_f = circular_velocity_hf(k, a_f) - beta_0_ = beta_0(V_0, V_f, inc_0, inc_f) + beta_0_ = _beta_0_hf(V_0, V_f, inc_0, inc_f) return V_0, V_f, beta_0_ -@jit -def delta_V(V_0, V_f, beta_0, inc_0, inc_f): +@hjit("f(f,f,f,f,f)") +def _delta_V_hf(V_0, V_f, beta_0, inc_0, inc_f): """Compute required increment of velocity.""" delta_i_f = abs(inc_f - inc_0) if delta_i_f == 0: return abs(V_f - V_0) - return V_0 * np.cos(beta_0) - V_0 * np.sin(beta_0) / np.tan( - np.pi / 2 * delta_i_f + beta_0 - ) + return V_0 * cos(beta_0) - V_0 * sin(beta_0) / tan(pi / 2 * delta_i_f + beta_0) + + +@hjit("Tuple([f,f])(f,f,f,f,f,f)") +def _extra_quantities_hf(k, a_0, a_f, inc_0, inc_f, f): + """Extra quantities given by the Edelbaum (a, i) model.""" + V_0, V_f, beta_0_ = _compute_parameters_hf(k, a_0, a_f, inc_0, inc_f) + delta_V = _delta_V_hf(V_0, V_f, beta_0_, inc_0, inc_f) + t_f_ = delta_V / f + + return delta_V, t_f_ + + +@hjit("f(f,f,f,f)") +def _beta_hf(t, V_0, f, beta_0): + """Compute yaw angle (β) as a function of time and the problem parameters.""" + return atan2(V_0 * sin(beta_0), V_0 * cos(beta_0) - f * t) -def change_a_inc(k, a_0, a_f, inc_0, inc_f, f): +def change_a_inc_hb(k, a_0, a_f, inc_0, inc_f, f): """Change semimajor axis and inclination. Guidance law from the Edelbaum/Kéchichian theory, optimal transfer between circular inclined orbits (a_0, i_0) --> (a_f, i_f), ecc = 0. @@ -87,25 +87,25 @@ def change_a_inc(k, a_0, a_f, inc_0, inc_f, f): References ---------- - * Edelbaum, T. N. "Propulsion Requirements for Controllable + * Edelbaum, T. N. "Propulsion Requirements delta_Vfor Controllable Satellites", 1961. * Kéchichian, J. A. "Reformulation of Edelbaum's Low-Thrust Transfer Problem Using Optimal Control Theory", 1997. """ - V_0, V_f, beta_0_ = compute_parameters(k, a_0, a_f, inc_0, inc_f) - - @jit - def a_d(t0, u_, k): - r = u_[:3] - v = u_[3:] + V_0, _, beta_0_ = _compute_parameters_hf(k, a_0, a_f, inc_0, inc_f) + @hjit("V(f,V,V,f)") + def a_d_hf(t0, rr, vv, k): # Change sign of beta with the out-of-plane velocity - beta_ = beta(t0, V_0, f, beta_0_) * np.sign(r[0] * (inc_f - inc_0)) - - t_ = v / norm_hf(array_to_V_hf(v)) - w_ = cross(r, v) / norm_hf(array_to_V_hf(cross(r, v))) - accel_v = f * (np.cos(beta_) * t_ + np.sin(beta_) * w_) + beta_ = _beta_hf(t0, V_0, f, beta_0_) * sign_hf(rr[0] * (inc_f - inc_0)) + + t_ = div_Vs_hf(vv, norm_hf(vv)) + crv = cross_VV_hf(rr, vv) + w_ = div_Vs_hf(crv, norm_hf(crv)) + accel_v = mul_Vs_hf( + add_VV_hf(mul_Vs_hf(t_, cos(beta_)), mul_Vs_hf(w_, sin(beta_))), f + ) return accel_v - delta_V, t_f = extra_quantities(k, a_0, a_f, inc_0, inc_f, f) - return a_d, delta_V, t_f + delta_V, t_f = _extra_quantities_hf(k, a_0, a_f, inc_0, inc_f, f) + return a_d_hf, delta_V, t_f diff --git a/src/hapsira/twobody/thrust/__init__.py b/src/hapsira/twobody/thrust/__init__.py index 924fcb10a..f91438807 100644 --- a/src/hapsira/twobody/thrust/__init__.py +++ b/src/hapsira/twobody/thrust/__init__.py @@ -1,9 +1,7 @@ -from hapsira.twobody.thrust.change_a_inc import change_a_inc -from hapsira.twobody.thrust.change_argp import change_argp -from hapsira.twobody.thrust.change_ecc_inc import change_ecc_inc -from hapsira.twobody.thrust.change_ecc_quasioptimal import ( - change_ecc_quasioptimal, -) +from .change_a_inc import change_a_inc +from .change_argp import change_argp +from .change_ecc_inc import change_ecc_inc +from .change_ecc_quasioptimal import change_ecc_quasioptimal __all__ = [ "change_a_inc", diff --git a/src/hapsira/twobody/thrust/change_a_inc.py b/src/hapsira/twobody/thrust/change_a_inc.py index fb391255f..c2e677299 100644 --- a/src/hapsira/twobody/thrust/change_a_inc.py +++ b/src/hapsira/twobody/thrust/change_a_inc.py @@ -1,8 +1,6 @@ from astropy import units as u -from hapsira.core.thrust.change_a_inc import ( - change_a_inc as change_a_inc_fast, -) +from hapsira.core.thrust.change_a_inc import change_a_inc_hb def change_a_inc(k, a_0, a_f, inc_0, inc_f, f): @@ -40,7 +38,7 @@ def change_a_inc(k, a_0, a_f, inc_0, inc_f, f): * Kéchichian, J. A. "Reformulation of Edelbaum's Low-Thrust Transfer Problem Using Optimal Control Theory", 1997. """ - a_d, delta_V, t_f = change_a_inc_fast( + a_d_hf, delta_V, t_f = change_a_inc_hb( k=k.to_value(u.km**3 / u.s**2), a_0=a_0.to_value(u.km), a_f=a_f.to_value(u.km), @@ -48,4 +46,8 @@ def change_a_inc(k, a_0, a_f, inc_0, inc_f, f): inc_f=inc_f.to_value(u.rad), f=f.to_value(u.km / u.s**2), ) - return a_d, delta_V, t_f * u.s + return ( + a_d_hf, + delta_V, + t_f * u.s, + ) # TODO delta_V is not a vector and does not carry a unit?? diff --git a/tests/tests_twobody/test_thrust.py b/tests/tests_twobody/test_thrust.py index 38a1b65f2..a40ba49e3 100644 --- a/tests/tests_twobody/test_thrust.py +++ b/tests/tests_twobody/test_thrust.py @@ -4,9 +4,10 @@ import pytest from hapsira.bodies import Earth +from hapsira.core.jit import array_to_V_hf from hapsira.core.propagation import func_twobody +from hapsira.core.thrust.change_a_inc import change_a_inc_hb from hapsira.core.thrust import ( - change_a_inc as change_a_inc_fast, change_argp as change_argp_fast, ) from hapsira.core.thrust.change_ecc_inc import beta as beta_change_ecc_inc @@ -42,7 +43,7 @@ def test_leo_geo_numerical_safe(inc_0): # Propagate orbit def f_leo_geo(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d(t0, u_, k) + ax, ay, az = a_d(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad @@ -66,7 +67,7 @@ def test_leo_geo_numerical_fast(inc_0): k = Earth.k.to(u.km**3 / u.s**2).value - a_d, _, t_f = change_a_inc_fast(k, a_0, a_f, inc_0, inc_f, f) + a_d_hf, _, t_f = change_a_inc_hb(k, a_0, a_f, inc_0, inc_f, f) # Retrieve r and v from initial orbit s0 = Orbit.circular(Earth, a_0 * u.km - Earth.R, inc_0 * u.rad) @@ -74,7 +75,7 @@ def test_leo_geo_numerical_fast(inc_0): # Propagate orbit def f_leo_geo(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d(t0, u_, k) + ax, ay, az = a_d_hf(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad @@ -354,7 +355,7 @@ def test_leo_geo_time_and_delta_v(inc_0, expected_t_f, expected_delta_V, rtol): k = Earth.k.to(u.km**3 / u.s**2).value inc_0 = np.radians(inc_0) # rad - _, delta_V, t_f = change_a_inc_fast(k, a_0, a_f, inc_0, inc_f, f) + _, delta_V, t_f = change_a_inc_hb(k, a_0, a_f, inc_0, inc_f, f) assert_allclose(delta_V, expected_delta_V, rtol=rtol) assert_allclose((t_f * u.s).to(u.day).value, expected_t_f, rtol=rtol) From 25c12e338d6dd2d3cb9bf4e87491c27327df31da Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 18:28:06 +0100 Subject: [PATCH 137/346] mask hf funcs --- src/hapsira/core/thrust/change_a_inc.py | 32 +++++++++++++++++++++---- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/src/hapsira/core/thrust/change_a_inc.py b/src/hapsira/core/thrust/change_a_inc.py index 202615c39..499887c98 100644 --- a/src/hapsira/core/thrust/change_a_inc.py +++ b/src/hapsira/core/thrust/change_a_inc.py @@ -1,6 +1,6 @@ from math import atan2, cos, pi, sin, tan -from ..jit import hjit +from ..jit import hjit, gjit from ..elements import circular_velocity_hf from ..math.linalg import add_VV_hf, cross_VV_hf, div_Vs_hf, mul_Vs_hf, norm_hf, sign_hf @@ -30,6 +30,15 @@ def _compute_parameters_hf(k, a_0, a_f, inc_0, inc_f): return V_0, V_f, beta_0_ +@gjit("void(f,f,f,f,f,f[:],f[:],f[:])", "(),(),(),(),()->(),(),()") +def _compute_parameters_gf(k, a_0, a_f, inc_0, inc_f, V_0, V_f, beta_0_): + """ + Vectorized compute_parameters + """ + + V_0[0], V_f[0], beta_0_[0] = _compute_parameters_hf(k, a_0, a_f, inc_0, inc_f) + + @hjit("f(f,f,f,f,f)") def _delta_V_hf(V_0, V_f, beta_0, inc_0, inc_f): """Compute required increment of velocity.""" @@ -49,6 +58,15 @@ def _extra_quantities_hf(k, a_0, a_f, inc_0, inc_f, f): return delta_V, t_f_ +@gjit("void(f,f,f,f,f,f,f[:],f[:])", "(),(),(),(),(),()->(),()") +def _extra_quantities_gf(k, a_0, a_f, inc_0, inc_f, f, delta_V, t_f_): + """ + Vectorized extra_quantities + """ + + delta_V[0], t_f_[0] = _extra_quantities_hf(k, a_0, a_f, inc_0, inc_f, f) + + @hjit("f(f,f,f,f)") def _beta_hf(t, V_0, f, beta_0): """Compute yaw angle (β) as a function of time and the problem parameters.""" @@ -78,7 +96,7 @@ def change_a_inc_hb(k, a_0, a_f, inc_0, inc_f, f): Returns ------- a_d : function - delta_V : numpy.ndarray + delta_V : float t_f : float Notes @@ -87,12 +105,14 @@ def change_a_inc_hb(k, a_0, a_f, inc_0, inc_f, f): References ---------- - * Edelbaum, T. N. "Propulsion Requirements delta_Vfor Controllable + * Edelbaum, T. N. "Propulsion Requirements delta_V for Controllable Satellites", 1961. * Kéchichian, J. A. "Reformulation of Edelbaum's Low-Thrust Transfer Problem Using Optimal Control Theory", 1997. """ - V_0, _, beta_0_ = _compute_parameters_hf(k, a_0, a_f, inc_0, inc_f) + V_0, _, beta_0_ = _compute_parameters_gf( # pylint: disable=E1120,E0633 + k, a_0, a_f, inc_0, inc_f + ) @hjit("V(f,V,V,f)") def a_d_hf(t0, rr, vv, k): @@ -107,5 +127,7 @@ def a_d_hf(t0, rr, vv, k): ) return accel_v - delta_V, t_f = _extra_quantities_hf(k, a_0, a_f, inc_0, inc_f, f) + delta_V, t_f = _extra_quantities_gf( # pylint: disable=E1120,E0633 + k, a_0, a_f, inc_0, inc_f, f + ) return a_d_hf, delta_V, t_f From fbe752c10788853ab58e33f114fad7bc839c77ec Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 18:28:13 +0100 Subject: [PATCH 138/346] fix typo --- src/hapsira/core/math/interpolate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/math/interpolate.py b/src/hapsira/core/math/interpolate.py index a70e654f5..a586286a9 100644 --- a/src/hapsira/core/math/interpolate.py +++ b/src/hapsira/core/math/interpolate.py @@ -15,7 +15,7 @@ def interp_hb(x: np.ndarray, y: np.ndarray) -> Callable: """ - Build compiled 1d-interpolator for 1D vectors + Build compiled 1d-interpolator for 3D vectors """ assert x.ndim == 1 From a083ab908fc34038cfed8169fc67cdc32af35a21 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 20:05:34 +0100 Subject: [PATCH 139/346] jit change_argp --- src/hapsira/core/thrust/change_argp.py | 66 ++++++++++++++--------- src/hapsira/twobody/thrust/change_argp.py | 8 +-- tests/tests_twobody/test_thrust.py | 14 +++-- 3 files changed, 50 insertions(+), 38 deletions(-) diff --git a/src/hapsira/core/thrust/change_argp.py b/src/hapsira/core/thrust/change_argp.py index dbd9a1429..1ba715b2a 100644 --- a/src/hapsira/core/thrust/change_argp.py +++ b/src/hapsira/core/thrust/change_argp.py @@ -1,33 +1,44 @@ -from numba import njit as jit -import numpy as np -from numpy import cross +from math import cos, pi, sin, sqrt -from hapsira.core.elements import circular_velocity_hf, rv2coe_hf, RV2COE_TOL +from ..elements import circular_velocity_hf, rv2coe_hf, RV2COE_TOL +from ..jit import hjit, gjit +from ..math.linalg import add_VV_hf, cross_VV_hf, div_Vs_hf, mul_Vs_hf, norm_hf, sign_hf -from ..jit import array_to_V_hf -from ..math.linalg import norm_hf +__all__ = [ + "change_argp_hb", +] -@jit -def delta_V(V, ecc, argp_0, argp_f, f, A): + +@hjit("f(f,f,f,f,f,f)") +def _delta_V_hf(V, ecc, argp_0, argp_f, f, A): """Compute required increment of velocity.""" delta_argp = argp_f - argp_0 return delta_argp / ( - 3 * np.sign(delta_argp) / 2 * np.sqrt(1 - ecc**2) / ecc / V + A / f + 3 * sign_hf(delta_argp) / 2 * sqrt(1 - ecc**2) / ecc / V + A / f ) -@jit -def extra_quantities(k, a, ecc, argp_0, argp_f, f, A=0.0): +@hjit("Tuple([f,f])(f,f,f,f,f,f,f)") +def _extra_quantities_hf(k, a, ecc, argp_0, argp_f, f, A): """Extra quantities given by the model.""" V = circular_velocity_hf(k, a) - delta_V_ = delta_V(V, ecc, argp_0, argp_f, f, A) + delta_V_ = _delta_V_hf(V, ecc, argp_0, argp_f, f, A) t_f_ = delta_V_ / f return delta_V_, t_f_ -def change_argp(k, a, ecc, argp_0, argp_f, f): +@gjit("void(f,f,f,f,f,f,f,f[:],f[:])", "(),(),(),(),(),(),()->(),()") +def _extra_quantities_gf(k, a, ecc, argp_0, argp_f, f, A, delta_V_, t_f_): + """ + Vectorized extra_quantities + """ + + delta_V_[0], t_f_[0] = _extra_quantities_hf(k, a, ecc, argp_0, argp_f, f, A) + + +def change_argp_hb(k, a, ecc, argp_0, argp_f, f): """Guidance law from the model. Thrust is aligned with an inertially fixed direction perpendicular to the semimajor axis of the orbit. @@ -50,24 +61,27 @@ def change_argp(k, a, ecc, argp_0, argp_f, f): Returns ------- a_d : function - delta_V : numpy.ndarray + delta_V : float t_f : float """ - @jit - def a_d(t0, u_, k): - r = u_[:3] - v = u_[3:] - nu = rv2coe_hf(k, array_to_V_hf(r), array_to_V_hf(v), RV2COE_TOL)[-1] + @hjit("V(f,V,V,f)") + def a_d_hf(t0, rr, vv, k): + nu = rv2coe_hf(k, rr, vv, RV2COE_TOL)[-1] - alpha_ = nu - np.pi / 2 + alpha_ = nu - pi / 2 - r_ = r / norm_hf(array_to_V_hf(r)) - w_ = cross(r, v) / norm_hf(array_to_V_hf(cross(r, v))) - s_ = cross(w_, r_) - accel_v = f * (np.cos(alpha_) * s_ + np.sin(alpha_) * r_) + r_ = div_Vs_hf(rr, norm_hf(rr)) + crv = cross_VV_hf(rr, vv) + w_ = div_Vs_hf(crv, norm_hf(crv)) + s_ = cross_VV_hf(w_, r_) + accel_v = mul_Vs_hf( + add_VV_hf(mul_Vs_hf(s_, cos(alpha_)), mul_Vs_hf(r_, sin(alpha_))), f + ) return accel_v - delta_V, t_f = extra_quantities(k, a, ecc, argp_0, argp_f, f, A=0.0) + delta_V, t_f = _extra_quantities_gf( # pylint: disable=E1120,E0633 + k, a, ecc, argp_0, argp_f, f, 0.0 + ) - return a_d, delta_V, t_f + return a_d_hf, delta_V, t_f diff --git a/src/hapsira/twobody/thrust/change_argp.py b/src/hapsira/twobody/thrust/change_argp.py index cc574234e..62eacbd6b 100644 --- a/src/hapsira/twobody/thrust/change_argp.py +++ b/src/hapsira/twobody/thrust/change_argp.py @@ -9,7 +9,7 @@ """ from astropy import units as u -from hapsira.core.thrust.change_argp import change_argp as change_a_inc_fast +from hapsira.core.thrust.change_argp import change_argp_hb def change_argp(k, a, ecc, argp_0, argp_f, f): @@ -36,13 +36,13 @@ def change_argp(k, a, ecc, argp_0, argp_f, f): ------- a_d, delta_V, t_f : tuple (function, ~astropy.units.quantity.Quantity, ~astropy.time.Time) """ - a_d, delta_V, t_f = change_a_inc_fast( + a_d_hf, delta_V, t_f = change_argp_hb( k=k.to_value(u.km**3 / u.s**2), a=a.to_value(u.km), - ecc=ecc, + ecc=ecc.to_value() if hasattr(ecc, "to_value") else ecc, argp_0=argp_0.to_value(u.rad), argp_f=argp_f.to_value(u.rad), f=f.to_value(u.km / u.s**2), ) - return a_d, delta_V, t_f * u.s + return a_d_hf, delta_V, t_f * u.s # delta_V is scalar, TODO add unit to it? diff --git a/tests/tests_twobody/test_thrust.py b/tests/tests_twobody/test_thrust.py index a40ba49e3..fd162a568 100644 --- a/tests/tests_twobody/test_thrust.py +++ b/tests/tests_twobody/test_thrust.py @@ -7,9 +7,7 @@ from hapsira.core.jit import array_to_V_hf from hapsira.core.propagation import func_twobody from hapsira.core.thrust.change_a_inc import change_a_inc_hb -from hapsira.core.thrust import ( - change_argp as change_argp_fast, -) +from hapsira.core.thrust.change_argp import change_argp_hb from hapsira.core.thrust.change_ecc_inc import beta as beta_change_ecc_inc from hapsira.twobody import Orbit from hapsira.twobody.propagation import CowellPropagator @@ -250,7 +248,7 @@ def test_soyuz_standard_gto_delta_v_fast(): k = Earth.k.to(u.km**3 / u.s**2).value - _, delta_V, t_f = change_argp_fast(k, a, ecc, argp_0, argp_f, f) + _, delta_V, t_f = change_argp_hb(k, a, ecc, argp_0, argp_f, f) expected_t_f = 12.0 # days, approximate expected_delta_V = 0.2489 # km / s @@ -272,7 +270,7 @@ def test_soyuz_standard_gto_numerical_safe(): k = Earth.k.to(u.km**3 / u.s**2) - a_d, _, t_f = change_argp(k, a, ecc, argp_0, argp_f, f) + a_d_hf, _, t_f = change_argp(k, a, ecc, argp_0, argp_f, f) # Retrieve r and v from initial orbit s0 = Orbit.from_classical( @@ -288,7 +286,7 @@ def test_soyuz_standard_gto_numerical_safe(): # Propagate orbit def f_soyuz(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d(t0, u_, k) + ax, ay, az = a_d_hf(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad @@ -310,7 +308,7 @@ def test_soyuz_standard_gto_numerical_fast(): k = Earth.k.to(u.km**3 / u.s**2).value - a_d, _, t_f = change_argp_fast(k, a, ecc, argp_0, argp_f, f) + a_d_hf, _, t_f = change_argp_hb(k, a, ecc, argp_0, argp_f, f) # Retrieve r and v from initial orbit s0 = Orbit.from_classical( @@ -326,7 +324,7 @@ def test_soyuz_standard_gto_numerical_fast(): # Propagate orbit def f_soyuz(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d(t0, u_, k) + ax, ay, az = a_d_hf(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad From efb3414934dbc5a72346981e4c5c87740cd511a5 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 20:05:53 +0100 Subject: [PATCH 140/346] rm exports --- src/hapsira/core/thrust/__init__.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/src/hapsira/core/thrust/__init__.py b/src/hapsira/core/thrust/__init__.py index 11a0ab926..e69de29bb 100644 --- a/src/hapsira/core/thrust/__init__.py +++ b/src/hapsira/core/thrust/__init__.py @@ -1,9 +0,0 @@ -# from .change_a_inc import change_a_inc_hb -from .change_argp import change_argp -from .change_ecc_inc import change_ecc_inc - -__all__ = [ - # "change_a_inc_hb", - "change_argp", - "change_ecc_inc", -] From 21e95189c7922d920f61fc54dbbe037b0936644b Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 21:04:55 +0100 Subject: [PATCH 141/346] jit change_ecc_inc --- ...tural-and-artificial-perturbations.myst.md | 7 +- src/hapsira/core/thrust/change_ecc_inc.py | 140 ++++++++++++------ src/hapsira/twobody/thrust/change_ecc_inc.py | 8 +- tests/tests_twobody/test_thrust.py | 10 +- 4 files changed, 108 insertions(+), 57 deletions(-) diff --git a/docs/source/examples/natural-and-artificial-perturbations.myst.md b/docs/source/examples/natural-and-artificial-perturbations.myst.md index 570cea410..4b14abeb9 100644 --- a/docs/source/examples/natural-and-artificial-perturbations.myst.md +++ b/docs/source/examples/natural-and-artificial-perturbations.myst.md @@ -262,14 +262,15 @@ orb0 = Orbit.from_classical( epoch=Time(0, format="jd", scale="tdb"), ) -a_d, _, t_f = change_ecc_inc(orb0, ecc_f, inc_f, f) +a_d_hf, _, t_f = change_ecc_inc(orb0, ecc_f, inc_f, f) def f(t0, state, k): du_kep = func_twobody(t0, state, k) - ax, ay, az = a_d( + ax, ay, az = a_d_hf( t0, - state, + array_to_V_hf(state[:3]), + array_to_V_hf(state[3:]), k, ) du_ad = np.array([0, 0, 0, ax, ay, az]) diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index 5218804d6..0cff79351 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -5,87 +5,141 @@ * Pollard, J. E. "Simplified Analysis of Low-Thrust Orbital Maneuvers", 2000. rv2coe """ -from numba import njit as jit -import numpy as np -from numpy import cross -from hapsira.core.elements import ( - circular_velocity_vf, - eccentricity_vector_gf, +from math import asin, atan, cos, pi, log, sin + +from numpy import array + +from ..elements import ( + circular_velocity_hf, + eccentricity_vector_hf, rv2coe_hf, RV2COE_TOL, ) +from ..jit import array_to_V_hf, hjit, vjit, gjit +from ..math.linalg import add_VV_hf, cross_VV_hf, div_Vs_hf, mul_Vs_hf, norm_hf, sign_hf + -from ..jit import array_to_V_hf -from ..math.linalg import norm_hf +__all__ = [ + "beta_hf", + "beta_vf", + "change_ecc_inc_hb", +] -@jit -def beta(ecc_0, ecc_f, inc_0, inc_f, argp): +@hjit("f(f,f,f,f,f)") +def beta_hf(ecc_0, ecc_f, inc_0, inc_f, argp): + """ + Scalar beta + """ # Note: "The argument of perigee will vary during the orbit transfer # due to the natural drift and because e may approach zero. # However, [the equation] still gives a good estimate of the desired # thrust angle." - return np.arctan( + return atan( abs( 3 - * np.pi + * pi * (inc_f - inc_0) / ( 4 - * np.cos(argp) + * cos(argp) * ( ecc_0 - ecc_f - + np.log((1 + ecc_f) * (-1 + ecc_0) / ((1 + ecc_0) * (-1 + ecc_f))) + + log((1 + ecc_f) * (-1 + ecc_0) / ((1 + ecc_0) * (-1 + ecc_f))) ) ) ) ) -@jit -def delta_V(V_0, ecc_0, ecc_f, beta_): - """Compute required increment of velocity.""" - return 2 * V_0 * np.abs(np.arcsin(ecc_0) - np.arcsin(ecc_f)) / (3 * np.cos(beta_)) +@vjit("f(f,f,f,f,f)") +def beta_vf(ecc_0, ecc_f, inc_0, inc_f, argp): + """ + Vectorized beta + """ + + return beta_hf(ecc_0, ecc_f, inc_0, inc_f, argp) + +@hjit("f(f,f,f,f)") +def _delta_V_hf(V_0, ecc_0, ecc_f, beta_): + """ + Compute required increment of velocity. + """ -@jit -def delta_t(delta_v, f): - """Compute required increment of velocity.""" + return 2 * V_0 * abs(asin(ecc_0) - asin(ecc_f)) / (3 * cos(beta_)) + + +@hjit("f(f,f)") +def _delta_t_hf(delta_v, f): + """ + Compute required increment of velocity. + """ return delta_v / f -def change_ecc_inc(k, a, ecc_0, ecc_f, inc_0, inc_f, argp, r, v, f): +@hjit("Tuple([V,f,f,f])(f,f,f,f,f,f,f,V,V,f)") +def _prepare_hf(k, a, ecc_0, ecc_f, inc_0, inc_f, argp, r, v, f): + """ + Vectorized prepare + """ + # We fix the inertial direction at the beginning if ecc_0 > 0.001: # Arbitrary tolerance - e_vec = eccentricity_vector_gf(k, r, v) # pylint: disable=E1120 - ref_vec = e_vec / ecc_0 + e_vec = eccentricity_vector_hf(k, r, v) + ref_vec = div_Vs_hf(e_vec, ecc_0) else: - ref_vec = r / norm_hf(array_to_V_hf(r)) + ref_vec = div_Vs_hf(r, norm_hf(r)) + + h_vec = cross_VV_hf(r, v) # Specific angular momentum vector + h_unit = div_Vs_hf(h_vec, norm_hf(h_vec)) + thrust_unit = mul_Vs_hf(cross_VV_hf(h_unit, ref_vec), sign_hf(ecc_f - ecc_0)) + + beta_0 = beta_hf(ecc_0, ecc_f, inc_0, inc_f, argp) - h_vec = cross(r, v) # Specific angular momentum vector - h_unit = h_vec / norm_hf(array_to_V_hf(h_vec)) - thrust_unit = cross(h_unit, ref_vec) * np.sign(ecc_f - ecc_0) + delta_v = _delta_V_hf(circular_velocity_hf(k, a), ecc_0, ecc_f, beta_0) + t_f = _delta_t_hf(delta_v, f) - beta_0 = beta(ecc_0, ecc_f, inc_0, inc_f, argp) + return thrust_unit, beta_0, delta_v, t_f - @jit - def a_d(t0, u_, k_): - r_ = u_[:3] - v_ = u_[3:] - nu = rv2coe_hf(k_, array_to_V_hf(r_), array_to_V_hf(v_), RV2COE_TOL)[-1] - beta_ = beta_0 * np.sign( - np.cos(nu) + +@gjit( + "void(f,f,f,f,f,f,f,f[:],f[:],f,f[:],f[:],f[:],f[:])", + "(),(),(),(),(),(),(),(n),(n),()->(n),(),(),()", +) +def _prepare_gf( + k, a, ecc_0, ecc_f, inc_0, inc_f, argp, r, v, f, thrust_unit, beta_0, delta_v, t_f +): + """ + Vectorized prepare + """ + + thrust_unit[:], beta_0[0], delta_v[0], t_f[0] = _prepare_hf( + k, a, ecc_0, ecc_f, inc_0, inc_f, argp, array_to_V_hf(r), array_to_V_hf(v), f + ) + + +def change_ecc_inc_hb(k, a, ecc_0, ecc_f, inc_0, inc_f, argp, r, v, f): + thrust_unit, beta_0, delta_v, t_f = _prepare_gf( # pylint: disable=E1120,E0633 + k, a, ecc_0, ecc_f, inc_0, inc_f, argp, array(r), array(v), f + ) + thrust_unit = tuple(thrust_unit) + + @hjit("V(f,V,V,f)") + def a_d_hf(t0, rr, vv, k_): + nu = rv2coe_hf(k_, rr, vv, RV2COE_TOL)[-1] + beta_ = beta_0 * sign_hf( + cos(nu) ) # The sign of ß reverses at minor axis crossings - w_ = (cross(r_, v_) / norm_hf(array_to_V_hf(cross(r_, v_)))) * np.sign( - inc_f - inc_0 + w_ = mul_Vs_hf( + cross_VV_hf(rr, vv), sign_hf(inc_f - inc_0) / norm_hf(cross_VV_hf(rr, vv)) + ) + accel_v = mul_Vs_hf( + add_VV_hf(mul_Vs_hf(thrust_unit, cos(beta_)), mul_Vs_hf(w_, sin(beta_))), f ) - accel_v = f * (np.cos(beta_) * thrust_unit + np.sin(beta_) * w_) return accel_v - delta_v = delta_V(circular_velocity_vf(k, a), ecc_0, ecc_f, beta_0) - t_f = delta_t(delta_v, f) - - return a_d, delta_v, t_f + return a_d_hf, delta_v, t_f diff --git a/src/hapsira/twobody/thrust/change_ecc_inc.py b/src/hapsira/twobody/thrust/change_ecc_inc.py index 4c3ca0c45..9c437beab 100644 --- a/src/hapsira/twobody/thrust/change_ecc_inc.py +++ b/src/hapsira/twobody/thrust/change_ecc_inc.py @@ -1,8 +1,6 @@ from astropy import units as u -from hapsira.core.thrust.change_ecc_inc import ( - change_ecc_inc as change_ecc_inc_fast, -) +from hapsira.core.thrust.change_ecc_inc import change_ecc_inc_hb def change_ecc_inc(orb_0, ecc_f, inc_f, f): @@ -28,7 +26,7 @@ def change_ecc_inc(orb_0, ecc_f, inc_f, f): * Pollard, J. E. "Simplified Analysis of Low-Thrust Orbital Maneuvers", 2000. """ r, v = orb_0.rv() - a_d, delta_V, t_f = change_ecc_inc_fast( + a_d_hf, delta_V, t_f = change_ecc_inc_hb( k=orb_0.attractor.k.to_value(u.km**3 / u.s**2), a=orb_0.a.to_value(u.km), ecc_0=orb_0.ecc.value, @@ -40,4 +38,4 @@ def change_ecc_inc(orb_0, ecc_f, inc_f, f): v=v.to_value(u.km / u.s), f=f.to_value(u.km / u.s**2), ) - return a_d, delta_V << (u.km / u.s), t_f << u.s + return a_d_hf, delta_V << (u.km / u.s), t_f << u.s diff --git a/tests/tests_twobody/test_thrust.py b/tests/tests_twobody/test_thrust.py index fd162a568..d18beec1d 100644 --- a/tests/tests_twobody/test_thrust.py +++ b/tests/tests_twobody/test_thrust.py @@ -8,7 +8,7 @@ from hapsira.core.propagation import func_twobody from hapsira.core.thrust.change_a_inc import change_a_inc_hb from hapsira.core.thrust.change_argp import change_argp_hb -from hapsira.core.thrust.change_ecc_inc import beta as beta_change_ecc_inc +from hapsira.core.thrust.change_ecc_inc import beta_vf as beta_change_ecc_inc from hapsira.twobody import Orbit from hapsira.twobody.propagation import CowellPropagator from hapsira.twobody.thrust import ( @@ -171,9 +171,7 @@ def test_geo_cases_beta_dnd_delta_v(ecc_0, inc_f, expected_beta, expected_delta_ nu=0 * u.deg, ) - beta = beta_change_ecc_inc( - ecc_0=ecc_0, ecc_f=ecc_f, inc_0=inc_0, inc_f=inc_f, argp=argp - ) + beta = beta_change_ecc_inc(ecc_0, ecc_f, inc_0, inc_f, argp) _, delta_V, _ = change_ecc_inc(orb_0=s0, ecc_f=ecc_f, inc_f=inc_f * u.rad, f=f) assert_allclose(delta_V.to_value(u.km / u.s), expected_delta_V, rtol=1e-2) @@ -198,12 +196,12 @@ def test_geo_cases_numerical(ecc_0, ecc_f): argp=argp * u.deg, nu=0 * u.deg, ) - a_d, _, t_f = change_ecc_inc(orb_0=s0, ecc_f=ecc_f, inc_f=inc_f, f=f) + a_d_hf, _, t_f = change_ecc_inc(orb_0=s0, ecc_f=ecc_f, inc_f=inc_f, f=f) # Propagate orbit def f_geo(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d(t0, u_, k) + ax, ay, az = a_d_hf(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad From f52f8b6f274813ad7376602a68e38c9923c6f9f4 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 21:46:47 +0100 Subject: [PATCH 142/346] jit change_ecc_quasioptimal --- .../core/thrust/change_ecc_quasioptimal.py | 91 ++++++++++++++++--- .../twobody/thrust/change_ecc_quasioptimal.py | 39 +++----- tests/tests_twobody/test_thrust.py | 4 +- 3 files changed, 97 insertions(+), 37 deletions(-) diff --git a/src/hapsira/core/thrust/change_ecc_quasioptimal.py b/src/hapsira/core/thrust/change_ecc_quasioptimal.py index d26895522..582ec9217 100644 --- a/src/hapsira/core/thrust/change_ecc_quasioptimal.py +++ b/src/hapsira/core/thrust/change_ecc_quasioptimal.py @@ -1,20 +1,89 @@ -from numba import njit as jit -import numpy as np +from math import asin -from hapsira.core.elements import circular_velocity_hf +from numpy import array +from ..elements import circular_velocity_hf +from ..jit import array_to_V_hf, hjit, gjit +from ..math.linalg import cross_VV_hf, div_Vs_hf, mul_Vs_hf, norm_hf, sign_hf -@jit -def delta_V(V_0, ecc_0, ecc_f): - """Compute required increment of velocity.""" - return 2 / 3 * V_0 * np.abs(np.arcsin(ecc_0) - np.arcsin(ecc_f)) +__all__ = [ + "change_ecc_quasioptimal_hb", +] -@jit -def extra_quantities(k, a, ecc_0, ecc_f, f): - """Extra quantities given by the model.""" +@hjit("f(f,f,f)") +def _delta_V_hf(V_0, ecc_0, ecc_f): + """ + Compute required increment of velocity. + """ + + return 2 / 3 * V_0 * abs(asin(ecc_0) - asin(ecc_f)) + + +@hjit("Tuple([f,f])(f,f,f,f,f)") +def _extra_quantities_hf(k, a, ecc_0, ecc_f, f): + """ + Extra quantities given by the model. + """ + V_0 = circular_velocity_hf(k, a) - delta_V_ = delta_V(V_0, ecc_0, ecc_f) + delta_V_ = _delta_V_hf(V_0, ecc_0, ecc_f) t_f_ = delta_V_ / f return delta_V_, t_f_ + + +@gjit("void(f,f,f,f,f,f[:],f[:])", "(),(),(),(),()->(),()") +def _extra_quantities_gf(k, a, ecc_0, ecc_f, f, delta_V_, t_f_): + """ + Vectorized extra_quantities + """ + + delta_V_[0], t_f_[0] = _extra_quantities_hf(k, a, ecc_0, ecc_f, f) + + +@hjit("V(f,f,f,f,V,V,V)") +def _prepare_hf(k, a, ecc_0, ecc_f, e_vec, h_vec, r): + """ + Scalar prepare + """ + + if ecc_0 > 0.001: # Arbitrary tolerance + ref_vec = div_Vs_hf(e_vec, ecc_0) + else: + ref_vec = div_Vs_hf(r, norm_hf(r)) + + h_unit = div_Vs_hf(h_vec, norm_hf(h_vec)) + thrust_unit = mul_Vs_hf(cross_VV_hf(h_unit, ref_vec), sign_hf(ecc_f - ecc_0)) + + return thrust_unit + + +@gjit("void(f,f,f,f,f[:],f[:],f[:],f[:])", "(),(),(),(),(n),(n),(n)->(n)") +def _prepare_gf(k, a, ecc_0, ecc_f, e_vec, h_vec, r, thrust_unit): + """ + Vectorized prepare + """ + + thrust_unit[:] = _prepare_hf( + k, a, ecc_0, ecc_f, array_to_V_hf(e_vec), array_to_V_hf(h_vec), array_to_V_hf(r) + ) + + +def change_ecc_quasioptimal_hb(k, a, ecc_0, ecc_f, e_vec, h_vec, r, f): + # We fix the inertial direction at the beginning + + thrust_unit = _prepare_gf( # pylint: disable=E1120,E0633 + k, a, ecc_0, array(ecc_f), array(e_vec), array(h_vec), r + ) + thrust_unit = tuple(thrust_unit) + + @hjit("V(f,V,V,f)") + def a_d_hf(t0, rr, vv, k): + accel_v = mul_Vs_hf(thrust_unit, f) + return accel_v + + delta_V, t_f = _extra_quantities_gf( # pylint: disable=E1120,E0633 + k, a, ecc_0, ecc_f, f + ) + return a_d_hf, delta_V, t_f diff --git a/src/hapsira/twobody/thrust/change_ecc_quasioptimal.py b/src/hapsira/twobody/thrust/change_ecc_quasioptimal.py index c7f13305b..842a28e74 100644 --- a/src/hapsira/twobody/thrust/change_ecc_quasioptimal.py +++ b/src/hapsira/twobody/thrust/change_ecc_quasioptimal.py @@ -7,12 +7,9 @@ """ from astropy import units as u -from numba import njit -import numpy as np -from numpy import cross -from hapsira.core.thrust.change_ecc_quasioptimal import extra_quantities -from hapsira.util import norm +from hapsira.core.jit import array_to_V_hf +from hapsira.core.thrust.change_ecc_quasioptimal import change_ecc_quasioptimal_hb def change_ecc_quasioptimal(orb_0, ecc_f, f): @@ -29,22 +26,16 @@ def change_ecc_quasioptimal(orb_0, ecc_f, f): f : float Magnitude of constant acceleration """ - # We fix the inertial direction at the beginning - k = orb_0.attractor.k.to(u.km**3 / u.s**2).value - a = orb_0.a.to(u.km).value - ecc_0 = orb_0.ecc.value - if ecc_0 > 0.001: # Arbitrary tolerance - ref_vec = orb_0.e_vec / ecc_0 - else: - ref_vec = orb_0.r / norm(orb_0.r) - - h_unit = orb_0.h_vec / norm(orb_0.h_vec) - thrust_unit = cross(h_unit, ref_vec) * np.sign(ecc_f - ecc_0) - - @njit - def a_d(t0, u_, k): - accel_v = f * thrust_unit - return accel_v - - delta_V, t_f = extra_quantities(k, a, ecc_0, ecc_f, f) - return a_d, delta_V, t_f + + a_d_hf, delta_V, t_f = change_ecc_quasioptimal_hb( + orb_0.attractor.k.to(u.km**3 / u.s**2).value, # k + orb_0.a.to(u.km).value, # a + orb_0.ecc.value, # ecc_0 + ecc_f, + array_to_V_hf(orb_0.e_vec), # e_vec, + array_to_V_hf(orb_0.h_vec), # h_vec, + array_to_V_hf(orb_0.r), # r + f, + ) + + return a_d_hf, delta_V, t_f diff --git a/tests/tests_twobody/test_thrust.py b/tests/tests_twobody/test_thrust.py index d18beec1d..bf59399e4 100644 --- a/tests/tests_twobody/test_thrust.py +++ b/tests/tests_twobody/test_thrust.py @@ -127,12 +127,12 @@ def test_sso_disposal_numerical(ecc_0, ecc_f): argp=0 * u.deg, nu=0 * u.deg, ) - a_d, _, t_f = change_ecc_quasioptimal(s0, ecc_f, f) + a_d_hf, _, t_f = change_ecc_quasioptimal(s0, ecc_f, f) # Propagate orbit def f_ss0_disposal(t0, u_, k): du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d(t0, u_, k) + ax, ay, az = a_d_hf(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) du_ad = np.array([0, 0, 0, ax, ay, az]) return du_kep + du_ad From 77a6ce36d5639c610e806ad368facdf8a31026b8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Jan 2024 22:04:07 +0100 Subject: [PATCH 143/346] no state type? two vectors instead? --- src/hapsira/core/jit.py | 7 ++++--- src/hapsira/core/propagation/cowell.py | 4 ++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 1232475e2..537969064 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -61,9 +61,10 @@ def _parse_signatures(signature: str, noreturn: bool = False) -> Union[str, List signature = signature.replace("M", "Tuple([V,V,V])") # matrix is a tuple of vectors signature = signature.replace("V", "Tuple([f,f,f])") # vector is a tuple of floats - signature = signature.replace( - "S", "Tuple([f,f,f,f,f,f])" - ) # state, two vectors, is a tuple of floats TODO remove, use vectors instead? + # signature = signature.replace( + # "S", "Tuple([f,f,f,f,f,f])" + # ) # state, two vectors, is a tuple of floats TODO remove, use vectors instead? + assert "S" not in signature signature = signature.replace( "F", "FunctionType" ) # TODO does not work for CUDA yet diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index ff0d820cc..3615d586c 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -5,11 +5,11 @@ from ..propagation.base import func_twobody -def cowell_jit(func): +def cowelljit(func): """ Wrapper for hjit to track funcs for cowell """ - compiled = hjit("S(f,S,f)")(func) + compiled = hjit("Tuple([V,V])(f,V,V,f)")(func) compiled.cowell = None # for debugging return compiled From 1dab907cd99f1b123aa272b164a8223f672b610f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 22 Jan 2024 09:34:14 +0100 Subject: [PATCH 144/346] isolate rkstep --- src/hapsira/core/math/ivp/_rk.py | 694 +------------------------- src/hapsira/core/math/ivp/_rkstep.py | 699 +++++++++++++++++++++++++++ 2 files changed, 700 insertions(+), 693 deletions(-) create mode 100644 src/hapsira/core/math/ivp/_rkstep.py diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 6dd47dff2..c0dd3ad16 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -5,7 +5,7 @@ from numba import jit from . import _dop853_coefficients as dop853_coefficients -from ._dop853_coefficients import A, B, C +from ._rkstep import rk_step, N_RV, N_STAGES __all__ = [ "EPS", @@ -22,15 +22,9 @@ MAX_FACTOR = 10 # Maximum allowed increase in a step size. INTERPOLATOR_POWER = 7 -N_RV = 6 -N_STAGES = 12 N_STAGES_EXTENDED = 16 ERROR_ESTIMATOR_ORDER = 7 -A = A[:N_STAGES, :N_STAGES] -B = B -C = C[:N_STAGES] - @jit(nopython=False) def norm(x: np.ndarray) -> float: @@ -38,692 +32,6 @@ def norm(x: np.ndarray) -> float: return np.linalg.norm(x) / x.size**0.5 -@jit(nopython=False) -def rk_step( - fun: Callable, - t: float, - y: np.ndarray, - f: np.ndarray, - h: float, - K: np.ndarray, - argk: float, -) -> Tuple[np.ndarray, np.ndarray]: - """Perform a single Runge-Kutta step. - - This function computes a prediction of an explicit Runge-Kutta method and - also estimates the error of a less accurate method. - - Notation for Butcher tableau is as in [1]_. - - Parameters - ---------- - fun : callable - Right-hand side of the system. - t : float - Current time. - y : ndarray, shape (n,) - Current state. - f : ndarray, shape (n,) - Current value of the derivative, i.e., ``fun(x, y)``. - h : float - Step to use. - K : ndarray, shape (n_stages + 1, n) - Storage array for putting RK stages here. Stages are stored in rows. - The last row is a linear combination of the previous rows with - coefficients - - Returns - ------- - y_new : ndarray, shape (n,) - Solution at t + h computed with a higher accuracy. - f_new : ndarray, shape (n,) - Derivative ``fun(t + h, y_new)``. - - Const - ----- - A : ndarray, shape (n_stages, n_stages) - Coefficients for combining previous RK stages to compute the next - stage. For explicit methods the coefficients at and above the main - diagonal are zeros. - B : ndarray, shape (n_stages,) - Coefficients for combining RK stages for computing the final - prediction. - C : ndarray, shape (n_stages,) - Coefficients for incrementing time for consecutive RK stages. - The value for the first stage is always zero. - - References - ---------- - .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential - Equations I: Nonstiff Problems", Sec. II.4. - """ - - assert y.shape == (N_RV,) - assert f.shape == (N_RV,) - assert K.shape == (N_STAGES + 1, N_RV) - - assert A.shape == (N_STAGES, N_STAGES) - assert B.shape == (N_STAGES,) - assert C.shape == (N_STAGES,) - - K[0] = f - - # for s, (a, c) in enumerate(zip(A[1:], C[1:]), start=1): - # dy = np.dot(K[:s].T, a[:s]) * h - # K[s] = fun(t + c * h, y + dy) - - # dy = np.dot(K[:1].T, A[1, :1]) * h - dy = np.array( - [ - (K[0][0] * A[1][0]) * h, - (K[0][1] * A[1][0]) * h, - (K[0][2] * A[1][0]) * h, - (K[0][3] * A[1][0]) * h, - (K[0][4] * A[1][0]) * h, - (K[0][5] * A[1][0]) * h, - ] - ) - K[1] = fun(t + C[1] * h, y + dy, argk) - - # dy = np.dot(K[:2].T, A[2, :2]) * h - dy = np.array( - [ - (K[0][0] * A[2][0] + K[1][0] * A[2][1]) * h, - (K[0][1] * A[2][0] + K[1][1] * A[2][1]) * h, - (K[0][2] * A[2][0] + K[1][2] * A[2][1]) * h, - (K[0][3] * A[2][0] + K[1][3] * A[2][1]) * h, - (K[0][4] * A[2][0] + K[1][4] * A[2][1]) * h, - (K[0][5] * A[2][0] + K[1][5] * A[2][1]) * h, - ] - ) - K[2] = fun(t + C[2] * h, y + dy, argk) - - # dy = np.dot(K[:3].T, A[3, :3]) * h - dy = np.array( - [ - (K[0][0] * A[3][0] + K[1][0] * A[3][1] + K[2][0] * A[3][2]) * h, - (K[0][1] * A[3][0] + K[1][1] * A[3][1] + K[2][1] * A[3][2]) * h, - (K[0][2] * A[3][0] + K[1][2] * A[3][1] + K[2][2] * A[3][2]) * h, - (K[0][3] * A[3][0] + K[1][3] * A[3][1] + K[2][3] * A[3][2]) * h, - (K[0][4] * A[3][0] + K[1][4] * A[3][1] + K[2][4] * A[3][2]) * h, - (K[0][5] * A[3][0] + K[1][5] * A[3][1] + K[2][5] * A[3][2]) * h, - ] - ) - - K[3] = fun(t + C[3] * h, y + dy, argk) - - # dy = np.dot(K[:4].T, A[4, :4]) * h - dy = np.array( - [ - ( - K[0][0] * A[4][0] - + K[1][0] * A[4][1] - + K[2][0] * A[4][2] - + K[3][0] * A[4][3] - ) - * h, - ( - K[0][1] * A[4][0] - + K[1][1] * A[4][1] - + K[2][1] * A[4][2] - + K[3][1] * A[4][3] - ) - * h, - ( - K[0][2] * A[4][0] - + K[1][2] * A[4][1] - + K[2][2] * A[4][2] - + K[3][2] * A[4][3] - ) - * h, - ( - K[0][3] * A[4][0] - + K[1][3] * A[4][1] - + K[2][3] * A[4][2] - + K[3][3] * A[4][3] - ) - * h, - ( - K[0][4] * A[4][0] - + K[1][4] * A[4][1] - + K[2][4] * A[4][2] - + K[3][4] * A[4][3] - ) - * h, - ( - K[0][5] * A[4][0] - + K[1][5] * A[4][1] - + K[2][5] * A[4][2] - + K[3][5] * A[4][3] - ) - * h, - ] - ) - K[4] = fun(t + C[4] * h, y + dy, argk) - - # dy = np.dot(K[:5].T, A[5, :5]) * h - dy = np.array( - [ - ( - K[0][0] * A[5][0] - + K[1][0] * A[5][1] - + K[2][0] * A[5][2] - + K[3][0] * A[5][3] - + K[4][0] * A[5][4] - ) - * h, - ( - K[0][1] * A[5][0] - + K[1][1] * A[5][1] - + K[2][1] * A[5][2] - + K[3][1] * A[5][3] - + K[4][1] * A[5][4] - ) - * h, - ( - K[0][2] * A[5][0] - + K[1][2] * A[5][1] - + K[2][2] * A[5][2] - + K[3][2] * A[5][3] - + K[4][2] * A[5][4] - ) - * h, - ( - K[0][3] * A[5][0] - + K[1][3] * A[5][1] - + K[2][3] * A[5][2] - + K[3][3] * A[5][3] - + K[4][3] * A[5][4] - ) - * h, - ( - K[0][4] * A[5][0] - + K[1][4] * A[5][1] - + K[2][4] * A[5][2] - + K[3][4] * A[5][3] - + K[4][4] * A[5][4] - ) - * h, - ( - K[0][5] * A[5][0] - + K[1][5] * A[5][1] - + K[2][5] * A[5][2] - + K[3][5] * A[5][3] - + K[4][5] * A[5][4] - ) - * h, - ] - ) - K[5] = fun(t + C[5] * h, y + dy, argk) - - # dy = np.dot(K[:6].T, A[6, :6]) * h - dy = np.array( - [ - ( - K[0][0] * A[6][0] - + K[1][0] * A[6][1] - + K[2][0] * A[6][2] - + K[3][0] * A[6][3] - + K[4][0] * A[6][4] - + K[5][0] * A[6][5] - ) - * h, - ( - K[0][1] * A[6][0] - + K[1][1] * A[6][1] - + K[2][1] * A[6][2] - + K[3][1] * A[6][3] - + K[4][1] * A[6][4] - + K[5][1] * A[6][5] - ) - * h, - ( - K[0][2] * A[6][0] - + K[1][2] * A[6][1] - + K[2][2] * A[6][2] - + K[3][2] * A[6][3] - + K[4][2] * A[6][4] - + K[5][2] * A[6][5] - ) - * h, - ( - K[0][3] * A[6][0] - + K[1][3] * A[6][1] - + K[2][3] * A[6][2] - + K[3][3] * A[6][3] - + K[4][3] * A[6][4] - + K[5][3] * A[6][5] - ) - * h, - ( - K[0][4] * A[6][0] - + K[1][4] * A[6][1] - + K[2][4] * A[6][2] - + K[3][4] * A[6][3] - + K[4][4] * A[6][4] - + K[5][4] * A[6][5] - ) - * h, - ( - K[0][5] * A[6][0] - + K[1][5] * A[6][1] - + K[2][5] * A[6][2] - + K[3][5] * A[6][3] - + K[4][5] * A[6][4] - + K[5][5] * A[6][5] - ) - * h, - ] - ) - K[6] = fun(t + C[6] * h, y + dy, argk) - - # dy = np.dot(K[:7].T, A[7, :7]) * h - dy = np.array( - [ - ( - K[0][0] * A[7][0] - + K[1][0] * A[7][1] - + K[2][0] * A[7][2] - + K[3][0] * A[7][3] - + K[4][0] * A[7][4] - + K[5][0] * A[7][5] - + K[6][0] * A[7][6] - ) - * h, - ( - K[0][1] * A[7][0] - + K[1][1] * A[7][1] - + K[2][1] * A[7][2] - + K[3][1] * A[7][3] - + K[4][1] * A[7][4] - + K[5][1] * A[7][5] - + K[6][1] * A[7][6] - ) - * h, - ( - K[0][2] * A[7][0] - + K[1][2] * A[7][1] - + K[2][2] * A[7][2] - + K[3][2] * A[7][3] - + K[4][2] * A[7][4] - + K[5][2] * A[7][5] - + K[6][2] * A[7][6] - ) - * h, - ( - K[0][3] * A[7][0] - + K[1][3] * A[7][1] - + K[2][3] * A[7][2] - + K[3][3] * A[7][3] - + K[4][3] * A[7][4] - + K[5][3] * A[7][5] - + K[6][3] * A[7][6] - ) - * h, - ( - K[0][4] * A[7][0] - + K[1][4] * A[7][1] - + K[2][4] * A[7][2] - + K[3][4] * A[7][3] - + K[4][4] * A[7][4] - + K[5][4] * A[7][5] - + K[6][4] * A[7][6] - ) - * h, - ( - K[0][5] * A[7][0] - + K[1][5] * A[7][1] - + K[2][5] * A[7][2] - + K[3][5] * A[7][3] - + K[4][5] * A[7][4] - + K[5][5] * A[7][5] - + K[6][5] * A[7][6] - ) - * h, - ] - ) - - K[7] = fun(t + C[7] * h, y + dy, argk) - - # dy = np.dot(K[:8].T, A[8, :8]) * h - dy = np.array( - [ - ( - K[0][0] * A[8][0] - + K[1][0] * A[8][1] - + K[2][0] * A[8][2] - + K[3][0] * A[8][3] - + K[4][0] * A[8][4] - + K[5][0] * A[8][5] - + K[6][0] * A[8][6] - + K[7][0] * A[8][7] - ) - * h, - ( - K[0][1] * A[8][0] - + K[1][1] * A[8][1] - + K[2][1] * A[8][2] - + K[3][1] * A[8][3] - + K[4][1] * A[8][4] - + K[5][1] * A[8][5] - + K[6][1] * A[8][6] - + K[7][1] * A[8][7] - ) - * h, - ( - K[0][2] * A[8][0] - + K[1][2] * A[8][1] - + K[2][2] * A[8][2] - + K[3][2] * A[8][3] - + K[4][2] * A[8][4] - + K[5][2] * A[8][5] - + K[6][2] * A[8][6] - + K[7][2] * A[8][7] - ) - * h, - ( - K[0][3] * A[8][0] - + K[1][3] * A[8][1] - + K[2][3] * A[8][2] - + K[3][3] * A[8][3] - + K[4][3] * A[8][4] - + K[5][3] * A[8][5] - + K[6][3] * A[8][6] - + K[7][3] * A[8][7] - ) - * h, - ( - K[0][4] * A[8][0] - + K[1][4] * A[8][1] - + K[2][4] * A[8][2] - + K[3][4] * A[8][3] - + K[4][4] * A[8][4] - + K[5][4] * A[8][5] - + K[6][4] * A[8][6] - + K[7][4] * A[8][7] - ) - * h, - ( - K[0][5] * A[8][0] - + K[1][5] * A[8][1] - + K[2][5] * A[8][2] - + K[3][5] * A[8][3] - + K[4][5] * A[8][4] - + K[5][5] * A[8][5] - + K[6][5] * A[8][6] - + K[7][5] * A[8][7] - ) - * h, - ] - ) - K[8] = fun(t + C[8] * h, y + dy, argk) - - # dy = np.dot(K[:9].T, A[9, :9]) * h - dy = np.array( - [ - ( - K[0][0] * A[9][0] - + K[1][0] * A[9][1] - + K[2][0] * A[9][2] - + K[3][0] * A[9][3] - + K[4][0] * A[9][4] - + K[5][0] * A[9][5] - + K[6][0] * A[9][6] - + K[7][0] * A[9][7] - + K[8][0] * A[9][8] - ) - * h, - ( - K[0][1] * A[9][0] - + K[1][1] * A[9][1] - + K[2][1] * A[9][2] - + K[3][1] * A[9][3] - + K[4][1] * A[9][4] - + K[5][1] * A[9][5] - + K[6][1] * A[9][6] - + K[7][1] * A[9][7] - + K[8][1] * A[9][8] - ) - * h, - ( - K[0][2] * A[9][0] - + K[1][2] * A[9][1] - + K[2][2] * A[9][2] - + K[3][2] * A[9][3] - + K[4][2] * A[9][4] - + K[5][2] * A[9][5] - + K[6][2] * A[9][6] - + K[7][2] * A[9][7] - + K[8][2] * A[9][8] - ) - * h, - ( - K[0][3] * A[9][0] - + K[1][3] * A[9][1] - + K[2][3] * A[9][2] - + K[3][3] * A[9][3] - + K[4][3] * A[9][4] - + K[5][3] * A[9][5] - + K[6][3] * A[9][6] - + K[7][3] * A[9][7] - + K[8][3] * A[9][8] - ) - * h, - ( - K[0][4] * A[9][0] - + K[1][4] * A[9][1] - + K[2][4] * A[9][2] - + K[3][4] * A[9][3] - + K[4][4] * A[9][4] - + K[5][4] * A[9][5] - + K[6][4] * A[9][6] - + K[7][4] * A[9][7] - + K[8][4] * A[9][8] - ) - * h, - ( - K[0][5] * A[9][0] - + K[1][5] * A[9][1] - + K[2][5] * A[9][2] - + K[3][5] * A[9][3] - + K[4][5] * A[9][4] - + K[5][5] * A[9][5] - + K[6][5] * A[9][6] - + K[7][5] * A[9][7] - + K[8][5] * A[9][8] - ) - * h, - ] - ) - K[9] = fun(t + C[9] * h, y + dy, argk) - - # dy = np.dot(K[:10].T, A[10, :10]) * h - dy = np.array( - [ - ( - K[0][0] * A[10][0] - + K[1][0] * A[10][1] - + K[2][0] * A[10][2] - + K[3][0] * A[10][3] - + K[4][0] * A[10][4] - + K[5][0] * A[10][5] - + K[6][0] * A[10][6] - + K[7][0] * A[10][7] - + K[8][0] * A[10][8] - + K[9][0] * A[10][9] - ) - * h, - ( - K[0][1] * A[10][0] - + K[1][1] * A[10][1] - + K[2][1] * A[10][2] - + K[3][1] * A[10][3] - + K[4][1] * A[10][4] - + K[5][1] * A[10][5] - + K[6][1] * A[10][6] - + K[7][1] * A[10][7] - + K[8][1] * A[10][8] - + K[9][1] * A[10][9] - ) - * h, - ( - K[0][2] * A[10][0] - + K[1][2] * A[10][1] - + K[2][2] * A[10][2] - + K[3][2] * A[10][3] - + K[4][2] * A[10][4] - + K[5][2] * A[10][5] - + K[6][2] * A[10][6] - + K[7][2] * A[10][7] - + K[8][2] * A[10][8] - + K[9][2] * A[10][9] - ) - * h, - ( - K[0][3] * A[10][0] - + K[1][3] * A[10][1] - + K[2][3] * A[10][2] - + K[3][3] * A[10][3] - + K[4][3] * A[10][4] - + K[5][3] * A[10][5] - + K[6][3] * A[10][6] - + K[7][3] * A[10][7] - + K[8][3] * A[10][8] - + K[9][3] * A[10][9] - ) - * h, - ( - K[0][4] * A[10][0] - + K[1][4] * A[10][1] - + K[2][4] * A[10][2] - + K[3][4] * A[10][3] - + K[4][4] * A[10][4] - + K[5][4] * A[10][5] - + K[6][4] * A[10][6] - + K[7][4] * A[10][7] - + K[8][4] * A[10][8] - + K[9][4] * A[10][9] - ) - * h, - ( - K[0][5] * A[10][0] - + K[1][5] * A[10][1] - + K[2][5] * A[10][2] - + K[3][5] * A[10][3] - + K[4][5] * A[10][4] - + K[5][5] * A[10][5] - + K[6][5] * A[10][6] - + K[7][5] * A[10][7] - + K[8][5] * A[10][8] - + K[9][5] * A[10][9] - ) - * h, - ] - ) - K[10] = fun(t + C[10] * h, y + dy, argk) - - # dy = np.dot(K[:11].T, A[11, :11]) * h - dy = np.array( - [ - ( - K[0][0] * A[11][0] - + K[1][0] * A[11][1] - + K[2][0] * A[11][2] - + K[3][0] * A[11][3] - + K[4][0] * A[11][4] - + K[5][0] * A[11][5] - + K[6][0] * A[11][6] - + K[7][0] * A[11][7] - + K[8][0] * A[11][8] - + K[9][0] * A[11][9] - + K[10][0] * A[11][10] - ) - * h, - ( - K[0][1] * A[11][0] - + K[1][1] * A[11][1] - + K[2][1] * A[11][2] - + K[3][1] * A[11][3] - + K[4][1] * A[11][4] - + K[5][1] * A[11][5] - + K[6][1] * A[11][6] - + K[7][1] * A[11][7] - + K[8][1] * A[11][8] - + K[9][1] * A[11][9] - + K[10][1] * A[11][10] - ) - * h, - ( - K[0][2] * A[11][0] - + K[1][2] * A[11][1] - + K[2][2] * A[11][2] - + K[3][2] * A[11][3] - + K[4][2] * A[11][4] - + K[5][2] * A[11][5] - + K[6][2] * A[11][6] - + K[7][2] * A[11][7] - + K[8][2] * A[11][8] - + K[9][2] * A[11][9] - + K[10][2] * A[11][10] - ) - * h, - ( - K[0][3] * A[11][0] - + K[1][3] * A[11][1] - + K[2][3] * A[11][2] - + K[3][3] * A[11][3] - + K[4][3] * A[11][4] - + K[5][3] * A[11][5] - + K[6][3] * A[11][6] - + K[7][3] * A[11][7] - + K[8][3] * A[11][8] - + K[9][3] * A[11][9] - + K[10][3] * A[11][10] - ) - * h, - ( - K[0][4] * A[11][0] - + K[1][4] * A[11][1] - + K[2][4] * A[11][2] - + K[3][4] * A[11][3] - + K[4][4] * A[11][4] - + K[5][4] * A[11][5] - + K[6][4] * A[11][6] - + K[7][4] * A[11][7] - + K[8][4] * A[11][8] - + K[9][4] * A[11][9] - + K[10][4] * A[11][10] - ) - * h, - ( - K[0][5] * A[11][0] - + K[1][5] * A[11][1] - + K[2][5] * A[11][2] - + K[3][5] * A[11][3] - + K[4][5] * A[11][4] - + K[5][5] * A[11][5] - + K[6][5] * A[11][6] - + K[7][5] * A[11][7] - + K[8][5] * A[11][8] - + K[9][5] * A[11][9] - + K[10][5] * A[11][10] - ) - * h, - ] - ) - K[11] = fun(t + C[11] * h, y + dy, argk) - - y_new = y + h * np.dot(K[:-1].T, B) - f_new = fun(t + h, y_new, argk) - - K[-1] = f_new - - assert y_new.shape == (N_RV,) - assert f_new.shape == (N_RV,) - - return y_new, f_new - - @jit(nopython=False) def select_initial_step( fun: Callable, diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py new file mode 100644 index 000000000..22eb4f22b --- /dev/null +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -0,0 +1,699 @@ +from typing import Callable, Tuple + +import numpy as np +from numba import jit + +from ._dop853_coefficients import A, B, C + +N_RV = 6 +N_STAGES = 12 + +A = A[:N_STAGES, :N_STAGES] +# B = B +C = C[:N_STAGES] + + +@jit(nopython=False) +def rk_step( + fun: Callable, + t: float, + y: np.ndarray, + f: np.ndarray, + h: float, + K: np.ndarray, + argk: float, +) -> Tuple[np.ndarray, np.ndarray]: + """Perform a single Runge-Kutta step. + + This function computes a prediction of an explicit Runge-Kutta method and + also estimates the error of a less accurate method. + + Notation for Butcher tableau is as in [1]_. + + Parameters + ---------- + fun : callable + Right-hand side of the system. + t : float + Current time. + y : ndarray, shape (n,) + Current state. + f : ndarray, shape (n,) + Current value of the derivative, i.e., ``fun(x, y)``. + h : float + Step to use. + K : ndarray, shape (n_stages + 1, n) + Storage array for putting RK stages here. Stages are stored in rows. + The last row is a linear combination of the previous rows with + coefficients + + Returns + ------- + y_new : ndarray, shape (n,) + Solution at t + h computed with a higher accuracy. + f_new : ndarray, shape (n,) + Derivative ``fun(t + h, y_new)``. + + Const + ----- + A : ndarray, shape (n_stages, n_stages) + Coefficients for combining previous RK stages to compute the next + stage. For explicit methods the coefficients at and above the main + diagonal are zeros. + B : ndarray, shape (n_stages,) + Coefficients for combining RK stages for computing the final + prediction. + C : ndarray, shape (n_stages,) + Coefficients for incrementing time for consecutive RK stages. + The value for the first stage is always zero. + + References + ---------- + .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential + Equations I: Nonstiff Problems", Sec. II.4. + """ + + assert y.shape == (N_RV,) + assert f.shape == (N_RV,) + assert K.shape == (N_STAGES + 1, N_RV) + + assert A.shape == (N_STAGES, N_STAGES) + assert B.shape == (N_STAGES,) + assert C.shape == (N_STAGES,) + + K[0] = f + + # for s, (a, c) in enumerate(zip(A[1:], C[1:]), start=1): + # dy = np.dot(K[:s].T, a[:s]) * h + # K[s] = fun(t + c * h, y + dy) + + # dy = np.dot(K[:1].T, A[1, :1]) * h + dy = np.array( + [ + (K[0][0] * A[1][0]) * h, + (K[0][1] * A[1][0]) * h, + (K[0][2] * A[1][0]) * h, + (K[0][3] * A[1][0]) * h, + (K[0][4] * A[1][0]) * h, + (K[0][5] * A[1][0]) * h, + ] + ) + K[1] = fun(t + C[1] * h, y + dy, argk) + + # dy = np.dot(K[:2].T, A[2, :2]) * h + dy = np.array( + [ + (K[0][0] * A[2][0] + K[1][0] * A[2][1]) * h, + (K[0][1] * A[2][0] + K[1][1] * A[2][1]) * h, + (K[0][2] * A[2][0] + K[1][2] * A[2][1]) * h, + (K[0][3] * A[2][0] + K[1][3] * A[2][1]) * h, + (K[0][4] * A[2][0] + K[1][4] * A[2][1]) * h, + (K[0][5] * A[2][0] + K[1][5] * A[2][1]) * h, + ] + ) + K[2] = fun(t + C[2] * h, y + dy, argk) + + # dy = np.dot(K[:3].T, A[3, :3]) * h + dy = np.array( + [ + (K[0][0] * A[3][0] + K[1][0] * A[3][1] + K[2][0] * A[3][2]) * h, + (K[0][1] * A[3][0] + K[1][1] * A[3][1] + K[2][1] * A[3][2]) * h, + (K[0][2] * A[3][0] + K[1][2] * A[3][1] + K[2][2] * A[3][2]) * h, + (K[0][3] * A[3][0] + K[1][3] * A[3][1] + K[2][3] * A[3][2]) * h, + (K[0][4] * A[3][0] + K[1][4] * A[3][1] + K[2][4] * A[3][2]) * h, + (K[0][5] * A[3][0] + K[1][5] * A[3][1] + K[2][5] * A[3][2]) * h, + ] + ) + + K[3] = fun(t + C[3] * h, y + dy, argk) + + # dy = np.dot(K[:4].T, A[4, :4]) * h + dy = np.array( + [ + ( + K[0][0] * A[4][0] + + K[1][0] * A[4][1] + + K[2][0] * A[4][2] + + K[3][0] * A[4][3] + ) + * h, + ( + K[0][1] * A[4][0] + + K[1][1] * A[4][1] + + K[2][1] * A[4][2] + + K[3][1] * A[4][3] + ) + * h, + ( + K[0][2] * A[4][0] + + K[1][2] * A[4][1] + + K[2][2] * A[4][2] + + K[3][2] * A[4][3] + ) + * h, + ( + K[0][3] * A[4][0] + + K[1][3] * A[4][1] + + K[2][3] * A[4][2] + + K[3][3] * A[4][3] + ) + * h, + ( + K[0][4] * A[4][0] + + K[1][4] * A[4][1] + + K[2][4] * A[4][2] + + K[3][4] * A[4][3] + ) + * h, + ( + K[0][5] * A[4][0] + + K[1][5] * A[4][1] + + K[2][5] * A[4][2] + + K[3][5] * A[4][3] + ) + * h, + ] + ) + K[4] = fun(t + C[4] * h, y + dy, argk) + + # dy = np.dot(K[:5].T, A[5, :5]) * h + dy = np.array( + [ + ( + K[0][0] * A[5][0] + + K[1][0] * A[5][1] + + K[2][0] * A[5][2] + + K[3][0] * A[5][3] + + K[4][0] * A[5][4] + ) + * h, + ( + K[0][1] * A[5][0] + + K[1][1] * A[5][1] + + K[2][1] * A[5][2] + + K[3][1] * A[5][3] + + K[4][1] * A[5][4] + ) + * h, + ( + K[0][2] * A[5][0] + + K[1][2] * A[5][1] + + K[2][2] * A[5][2] + + K[3][2] * A[5][3] + + K[4][2] * A[5][4] + ) + * h, + ( + K[0][3] * A[5][0] + + K[1][3] * A[5][1] + + K[2][3] * A[5][2] + + K[3][3] * A[5][3] + + K[4][3] * A[5][4] + ) + * h, + ( + K[0][4] * A[5][0] + + K[1][4] * A[5][1] + + K[2][4] * A[5][2] + + K[3][4] * A[5][3] + + K[4][4] * A[5][4] + ) + * h, + ( + K[0][5] * A[5][0] + + K[1][5] * A[5][1] + + K[2][5] * A[5][2] + + K[3][5] * A[5][3] + + K[4][5] * A[5][4] + ) + * h, + ] + ) + K[5] = fun(t + C[5] * h, y + dy, argk) + + # dy = np.dot(K[:6].T, A[6, :6]) * h + dy = np.array( + [ + ( + K[0][0] * A[6][0] + + K[1][0] * A[6][1] + + K[2][0] * A[6][2] + + K[3][0] * A[6][3] + + K[4][0] * A[6][4] + + K[5][0] * A[6][5] + ) + * h, + ( + K[0][1] * A[6][0] + + K[1][1] * A[6][1] + + K[2][1] * A[6][2] + + K[3][1] * A[6][3] + + K[4][1] * A[6][4] + + K[5][1] * A[6][5] + ) + * h, + ( + K[0][2] * A[6][0] + + K[1][2] * A[6][1] + + K[2][2] * A[6][2] + + K[3][2] * A[6][3] + + K[4][2] * A[6][4] + + K[5][2] * A[6][5] + ) + * h, + ( + K[0][3] * A[6][0] + + K[1][3] * A[6][1] + + K[2][3] * A[6][2] + + K[3][3] * A[6][3] + + K[4][3] * A[6][4] + + K[5][3] * A[6][5] + ) + * h, + ( + K[0][4] * A[6][0] + + K[1][4] * A[6][1] + + K[2][4] * A[6][2] + + K[3][4] * A[6][3] + + K[4][4] * A[6][4] + + K[5][4] * A[6][5] + ) + * h, + ( + K[0][5] * A[6][0] + + K[1][5] * A[6][1] + + K[2][5] * A[6][2] + + K[3][5] * A[6][3] + + K[4][5] * A[6][4] + + K[5][5] * A[6][5] + ) + * h, + ] + ) + K[6] = fun(t + C[6] * h, y + dy, argk) + + # dy = np.dot(K[:7].T, A[7, :7]) * h + dy = np.array( + [ + ( + K[0][0] * A[7][0] + + K[1][0] * A[7][1] + + K[2][0] * A[7][2] + + K[3][0] * A[7][3] + + K[4][0] * A[7][4] + + K[5][0] * A[7][5] + + K[6][0] * A[7][6] + ) + * h, + ( + K[0][1] * A[7][0] + + K[1][1] * A[7][1] + + K[2][1] * A[7][2] + + K[3][1] * A[7][3] + + K[4][1] * A[7][4] + + K[5][1] * A[7][5] + + K[6][1] * A[7][6] + ) + * h, + ( + K[0][2] * A[7][0] + + K[1][2] * A[7][1] + + K[2][2] * A[7][2] + + K[3][2] * A[7][3] + + K[4][2] * A[7][4] + + K[5][2] * A[7][5] + + K[6][2] * A[7][6] + ) + * h, + ( + K[0][3] * A[7][0] + + K[1][3] * A[7][1] + + K[2][3] * A[7][2] + + K[3][3] * A[7][3] + + K[4][3] * A[7][4] + + K[5][3] * A[7][5] + + K[6][3] * A[7][6] + ) + * h, + ( + K[0][4] * A[7][0] + + K[1][4] * A[7][1] + + K[2][4] * A[7][2] + + K[3][4] * A[7][3] + + K[4][4] * A[7][4] + + K[5][4] * A[7][5] + + K[6][4] * A[7][6] + ) + * h, + ( + K[0][5] * A[7][0] + + K[1][5] * A[7][1] + + K[2][5] * A[7][2] + + K[3][5] * A[7][3] + + K[4][5] * A[7][4] + + K[5][5] * A[7][5] + + K[6][5] * A[7][6] + ) + * h, + ] + ) + + K[7] = fun(t + C[7] * h, y + dy, argk) + + # dy = np.dot(K[:8].T, A[8, :8]) * h + dy = np.array( + [ + ( + K[0][0] * A[8][0] + + K[1][0] * A[8][1] + + K[2][0] * A[8][2] + + K[3][0] * A[8][3] + + K[4][0] * A[8][4] + + K[5][0] * A[8][5] + + K[6][0] * A[8][6] + + K[7][0] * A[8][7] + ) + * h, + ( + K[0][1] * A[8][0] + + K[1][1] * A[8][1] + + K[2][1] * A[8][2] + + K[3][1] * A[8][3] + + K[4][1] * A[8][4] + + K[5][1] * A[8][5] + + K[6][1] * A[8][6] + + K[7][1] * A[8][7] + ) + * h, + ( + K[0][2] * A[8][0] + + K[1][2] * A[8][1] + + K[2][2] * A[8][2] + + K[3][2] * A[8][3] + + K[4][2] * A[8][4] + + K[5][2] * A[8][5] + + K[6][2] * A[8][6] + + K[7][2] * A[8][7] + ) + * h, + ( + K[0][3] * A[8][0] + + K[1][3] * A[8][1] + + K[2][3] * A[8][2] + + K[3][3] * A[8][3] + + K[4][3] * A[8][4] + + K[5][3] * A[8][5] + + K[6][3] * A[8][6] + + K[7][3] * A[8][7] + ) + * h, + ( + K[0][4] * A[8][0] + + K[1][4] * A[8][1] + + K[2][4] * A[8][2] + + K[3][4] * A[8][3] + + K[4][4] * A[8][4] + + K[5][4] * A[8][5] + + K[6][4] * A[8][6] + + K[7][4] * A[8][7] + ) + * h, + ( + K[0][5] * A[8][0] + + K[1][5] * A[8][1] + + K[2][5] * A[8][2] + + K[3][5] * A[8][3] + + K[4][5] * A[8][4] + + K[5][5] * A[8][5] + + K[6][5] * A[8][6] + + K[7][5] * A[8][7] + ) + * h, + ] + ) + K[8] = fun(t + C[8] * h, y + dy, argk) + + # dy = np.dot(K[:9].T, A[9, :9]) * h + dy = np.array( + [ + ( + K[0][0] * A[9][0] + + K[1][0] * A[9][1] + + K[2][0] * A[9][2] + + K[3][0] * A[9][3] + + K[4][0] * A[9][4] + + K[5][0] * A[9][5] + + K[6][0] * A[9][6] + + K[7][0] * A[9][7] + + K[8][0] * A[9][8] + ) + * h, + ( + K[0][1] * A[9][0] + + K[1][1] * A[9][1] + + K[2][1] * A[9][2] + + K[3][1] * A[9][3] + + K[4][1] * A[9][4] + + K[5][1] * A[9][5] + + K[6][1] * A[9][6] + + K[7][1] * A[9][7] + + K[8][1] * A[9][8] + ) + * h, + ( + K[0][2] * A[9][0] + + K[1][2] * A[9][1] + + K[2][2] * A[9][2] + + K[3][2] * A[9][3] + + K[4][2] * A[9][4] + + K[5][2] * A[9][5] + + K[6][2] * A[9][6] + + K[7][2] * A[9][7] + + K[8][2] * A[9][8] + ) + * h, + ( + K[0][3] * A[9][0] + + K[1][3] * A[9][1] + + K[2][3] * A[9][2] + + K[3][3] * A[9][3] + + K[4][3] * A[9][4] + + K[5][3] * A[9][5] + + K[6][3] * A[9][6] + + K[7][3] * A[9][7] + + K[8][3] * A[9][8] + ) + * h, + ( + K[0][4] * A[9][0] + + K[1][4] * A[9][1] + + K[2][4] * A[9][2] + + K[3][4] * A[9][3] + + K[4][4] * A[9][4] + + K[5][4] * A[9][5] + + K[6][4] * A[9][6] + + K[7][4] * A[9][7] + + K[8][4] * A[9][8] + ) + * h, + ( + K[0][5] * A[9][0] + + K[1][5] * A[9][1] + + K[2][5] * A[9][2] + + K[3][5] * A[9][3] + + K[4][5] * A[9][4] + + K[5][5] * A[9][5] + + K[6][5] * A[9][6] + + K[7][5] * A[9][7] + + K[8][5] * A[9][8] + ) + * h, + ] + ) + K[9] = fun(t + C[9] * h, y + dy, argk) + + # dy = np.dot(K[:10].T, A[10, :10]) * h + dy = np.array( + [ + ( + K[0][0] * A[10][0] + + K[1][0] * A[10][1] + + K[2][0] * A[10][2] + + K[3][0] * A[10][3] + + K[4][0] * A[10][4] + + K[5][0] * A[10][5] + + K[6][0] * A[10][6] + + K[7][0] * A[10][7] + + K[8][0] * A[10][8] + + K[9][0] * A[10][9] + ) + * h, + ( + K[0][1] * A[10][0] + + K[1][1] * A[10][1] + + K[2][1] * A[10][2] + + K[3][1] * A[10][3] + + K[4][1] * A[10][4] + + K[5][1] * A[10][5] + + K[6][1] * A[10][6] + + K[7][1] * A[10][7] + + K[8][1] * A[10][8] + + K[9][1] * A[10][9] + ) + * h, + ( + K[0][2] * A[10][0] + + K[1][2] * A[10][1] + + K[2][2] * A[10][2] + + K[3][2] * A[10][3] + + K[4][2] * A[10][4] + + K[5][2] * A[10][5] + + K[6][2] * A[10][6] + + K[7][2] * A[10][7] + + K[8][2] * A[10][8] + + K[9][2] * A[10][9] + ) + * h, + ( + K[0][3] * A[10][0] + + K[1][3] * A[10][1] + + K[2][3] * A[10][2] + + K[3][3] * A[10][3] + + K[4][3] * A[10][4] + + K[5][3] * A[10][5] + + K[6][3] * A[10][6] + + K[7][3] * A[10][7] + + K[8][3] * A[10][8] + + K[9][3] * A[10][9] + ) + * h, + ( + K[0][4] * A[10][0] + + K[1][4] * A[10][1] + + K[2][4] * A[10][2] + + K[3][4] * A[10][3] + + K[4][4] * A[10][4] + + K[5][4] * A[10][5] + + K[6][4] * A[10][6] + + K[7][4] * A[10][7] + + K[8][4] * A[10][8] + + K[9][4] * A[10][9] + ) + * h, + ( + K[0][5] * A[10][0] + + K[1][5] * A[10][1] + + K[2][5] * A[10][2] + + K[3][5] * A[10][3] + + K[4][5] * A[10][4] + + K[5][5] * A[10][5] + + K[6][5] * A[10][6] + + K[7][5] * A[10][7] + + K[8][5] * A[10][8] + + K[9][5] * A[10][9] + ) + * h, + ] + ) + K[10] = fun(t + C[10] * h, y + dy, argk) + + # dy = np.dot(K[:11].T, A[11, :11]) * h + dy = np.array( + [ + ( + K[0][0] * A[11][0] + + K[1][0] * A[11][1] + + K[2][0] * A[11][2] + + K[3][0] * A[11][3] + + K[4][0] * A[11][4] + + K[5][0] * A[11][5] + + K[6][0] * A[11][6] + + K[7][0] * A[11][7] + + K[8][0] * A[11][8] + + K[9][0] * A[11][9] + + K[10][0] * A[11][10] + ) + * h, + ( + K[0][1] * A[11][0] + + K[1][1] * A[11][1] + + K[2][1] * A[11][2] + + K[3][1] * A[11][3] + + K[4][1] * A[11][4] + + K[5][1] * A[11][5] + + K[6][1] * A[11][6] + + K[7][1] * A[11][7] + + K[8][1] * A[11][8] + + K[9][1] * A[11][9] + + K[10][1] * A[11][10] + ) + * h, + ( + K[0][2] * A[11][0] + + K[1][2] * A[11][1] + + K[2][2] * A[11][2] + + K[3][2] * A[11][3] + + K[4][2] * A[11][4] + + K[5][2] * A[11][5] + + K[6][2] * A[11][6] + + K[7][2] * A[11][7] + + K[8][2] * A[11][8] + + K[9][2] * A[11][9] + + K[10][2] * A[11][10] + ) + * h, + ( + K[0][3] * A[11][0] + + K[1][3] * A[11][1] + + K[2][3] * A[11][2] + + K[3][3] * A[11][3] + + K[4][3] * A[11][4] + + K[5][3] * A[11][5] + + K[6][3] * A[11][6] + + K[7][3] * A[11][7] + + K[8][3] * A[11][8] + + K[9][3] * A[11][9] + + K[10][3] * A[11][10] + ) + * h, + ( + K[0][4] * A[11][0] + + K[1][4] * A[11][1] + + K[2][4] * A[11][2] + + K[3][4] * A[11][3] + + K[4][4] * A[11][4] + + K[5][4] * A[11][5] + + K[6][4] * A[11][6] + + K[7][4] * A[11][7] + + K[8][4] * A[11][8] + + K[9][4] * A[11][9] + + K[10][4] * A[11][10] + ) + * h, + ( + K[0][5] * A[11][0] + + K[1][5] * A[11][1] + + K[2][5] * A[11][2] + + K[3][5] * A[11][3] + + K[4][5] * A[11][4] + + K[5][5] * A[11][5] + + K[6][5] * A[11][6] + + K[7][5] * A[11][7] + + K[8][5] * A[11][8] + + K[9][5] * A[11][9] + + K[10][5] * A[11][10] + ) + * h, + ] + ) + K[11] = fun(t + C[11] * h, y + dy, argk) + + y_new = y + h * np.dot(K[:-1].T, B) + f_new = fun(t + h, y_new, argk) + + K[-1] = f_new + + assert y_new.shape == (N_RV,) + assert f_new.shape == (N_RV,) + + return y_new, f_new From 6ece659583b5135ac86cb5160a5c02867ca93835 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 22 Jan 2024 09:48:42 +0100 Subject: [PATCH 145/346] run on tuples --- src/hapsira/core/math/ivp/_rkstep.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index 22eb4f22b..dd86776ae 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -5,12 +5,18 @@ from ._dop853_coefficients import A, B, C +__all__ = [ + "rk_step", + "N_RV", + "N_STAGES", +] + N_RV = 6 N_STAGES = 12 -A = A[:N_STAGES, :N_STAGES] +A = tuple(tuple(line) for line in A[:N_STAGES, :N_STAGES]) # B = B -C = C[:N_STAGES] +C = tuple(C[:N_STAGES]) @jit(nopython=False) @@ -77,9 +83,9 @@ def rk_step( assert f.shape == (N_RV,) assert K.shape == (N_STAGES + 1, N_RV) - assert A.shape == (N_STAGES, N_STAGES) + # assert A.shape == (N_STAGES, N_STAGES) assert B.shape == (N_STAGES,) - assert C.shape == (N_STAGES,) + # assert C.shape == (N_STAGES,) K[0] = f From 000faa70d62e553d2d9c340764ddcb1950b2a9fa Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 22 Jan 2024 11:11:47 +0100 Subject: [PATCH 146/346] run rkstep on tuples without side effects --- src/hapsira/core/math/ivp/_rk.py | 4 +- src/hapsira/core/math/ivp/_rkstep.py | 886 ++++++++++++++------------- 2 files changed, 448 insertions(+), 442 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index c0dd3ad16..428b78042 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -420,7 +420,9 @@ def _step_impl(self): h = t_new - t h_abs = np.abs(h) - y_new, f_new = rk_step(self.fun, t, y, self.f, h, self.K, self.argk) + y_new, f_new, K_new = rk_step(self.fun, t, y, self.f, h, self.argk) + self.K[: N_STAGES + 1, :N_RV] = np.array([K_new]) + scale = atol + np.maximum(np.abs(y), np.abs(y_new)) * rtol error_norm = self._estimate_error_norm(self.K, h, scale) diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index dd86776ae..4dafac978 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -26,7 +26,6 @@ def rk_step( y: np.ndarray, f: np.ndarray, h: float, - K: np.ndarray, argk: float, ) -> Tuple[np.ndarray, np.ndarray]: """Perform a single Runge-Kutta step. @@ -48,10 +47,6 @@ def rk_step( Current value of the derivative, i.e., ``fun(x, y)``. h : float Step to use. - K : ndarray, shape (n_stages + 1, n) - Storage array for putting RK stages here. Stages are stored in rows. - The last row is a linear combination of the previous rows with - coefficients Returns ------- @@ -59,6 +54,10 @@ def rk_step( Solution at t + h computed with a higher accuracy. f_new : ndarray, shape (n,) Derivative ``fun(t + h, y_new)``. + K : ndarray, shape (n_stages + 1, n) + Storage array for putting RK stages here. Stages are stored in rows. + The last row is a linear combination of the previous rows with + coefficients Const ----- @@ -81,13 +80,13 @@ def rk_step( assert y.shape == (N_RV,) assert f.shape == (N_RV,) - assert K.shape == (N_STAGES + 1, N_RV) + # assert K.shape == (N_STAGES + 1, N_RV) # assert A.shape == (N_STAGES, N_STAGES) assert B.shape == (N_STAGES,) # assert C.shape == (N_STAGES,) - K[0] = f + K00 = f # for s, (a, c) in enumerate(zip(A[1:], C[1:]), start=1): # dy = np.dot(K[:s].T, a[:s]) * h @@ -96,610 +95,615 @@ def rk_step( # dy = np.dot(K[:1].T, A[1, :1]) * h dy = np.array( [ - (K[0][0] * A[1][0]) * h, - (K[0][1] * A[1][0]) * h, - (K[0][2] * A[1][0]) * h, - (K[0][3] * A[1][0]) * h, - (K[0][4] * A[1][0]) * h, - (K[0][5] * A[1][0]) * h, + (K00[0] * A[1][0]) * h, + (K00[1] * A[1][0]) * h, + (K00[2] * A[1][0]) * h, + (K00[3] * A[1][0]) * h, + (K00[4] * A[1][0]) * h, + (K00[5] * A[1][0]) * h, ] ) - K[1] = fun(t + C[1] * h, y + dy, argk) + K01 = fun(t + C[1] * h, y + dy, argk) # dy = np.dot(K[:2].T, A[2, :2]) * h dy = np.array( [ - (K[0][0] * A[2][0] + K[1][0] * A[2][1]) * h, - (K[0][1] * A[2][0] + K[1][1] * A[2][1]) * h, - (K[0][2] * A[2][0] + K[1][2] * A[2][1]) * h, - (K[0][3] * A[2][0] + K[1][3] * A[2][1]) * h, - (K[0][4] * A[2][0] + K[1][4] * A[2][1]) * h, - (K[0][5] * A[2][0] + K[1][5] * A[2][1]) * h, + (K00[0] * A[2][0] + K01[0] * A[2][1]) * h, + (K00[1] * A[2][0] + K01[1] * A[2][1]) * h, + (K00[2] * A[2][0] + K01[2] * A[2][1]) * h, + (K00[3] * A[2][0] + K01[3] * A[2][1]) * h, + (K00[4] * A[2][0] + K01[4] * A[2][1]) * h, + (K00[5] * A[2][0] + K01[5] * A[2][1]) * h, ] ) - K[2] = fun(t + C[2] * h, y + dy, argk) + K02 = fun(t + C[2] * h, y + dy, argk) # dy = np.dot(K[:3].T, A[3, :3]) * h dy = np.array( [ - (K[0][0] * A[3][0] + K[1][0] * A[3][1] + K[2][0] * A[3][2]) * h, - (K[0][1] * A[3][0] + K[1][1] * A[3][1] + K[2][1] * A[3][2]) * h, - (K[0][2] * A[3][0] + K[1][2] * A[3][1] + K[2][2] * A[3][2]) * h, - (K[0][3] * A[3][0] + K[1][3] * A[3][1] + K[2][3] * A[3][2]) * h, - (K[0][4] * A[3][0] + K[1][4] * A[3][1] + K[2][4] * A[3][2]) * h, - (K[0][5] * A[3][0] + K[1][5] * A[3][1] + K[2][5] * A[3][2]) * h, + (K00[0] * A[3][0] + K01[0] * A[3][1] + K02[0] * A[3][2]) * h, + (K00[1] * A[3][0] + K01[1] * A[3][1] + K02[1] * A[3][2]) * h, + (K00[2] * A[3][0] + K01[2] * A[3][1] + K02[2] * A[3][2]) * h, + (K00[3] * A[3][0] + K01[3] * A[3][1] + K02[3] * A[3][2]) * h, + (K00[4] * A[3][0] + K01[4] * A[3][1] + K02[4] * A[3][2]) * h, + (K00[5] * A[3][0] + K01[5] * A[3][1] + K02[5] * A[3][2]) * h, ] ) - K[3] = fun(t + C[3] * h, y + dy, argk) + K03 = fun(t + C[3] * h, y + dy, argk) # dy = np.dot(K[:4].T, A[4, :4]) * h dy = np.array( [ - ( - K[0][0] * A[4][0] - + K[1][0] * A[4][1] - + K[2][0] * A[4][2] - + K[3][0] * A[4][3] - ) + (K00[0] * A[4][0] + K01[0] * A[4][1] + K02[0] * A[4][2] + K03[0] * A[4][3]) * h, - ( - K[0][1] * A[4][0] - + K[1][1] * A[4][1] - + K[2][1] * A[4][2] - + K[3][1] * A[4][3] - ) + (K00[1] * A[4][0] + K01[1] * A[4][1] + K02[1] * A[4][2] + K03[1] * A[4][3]) * h, - ( - K[0][2] * A[4][0] - + K[1][2] * A[4][1] - + K[2][2] * A[4][2] - + K[3][2] * A[4][3] - ) + (K00[2] * A[4][0] + K01[2] * A[4][1] + K02[2] * A[4][2] + K03[2] * A[4][3]) * h, - ( - K[0][3] * A[4][0] - + K[1][3] * A[4][1] - + K[2][3] * A[4][2] - + K[3][3] * A[4][3] - ) + (K00[3] * A[4][0] + K01[3] * A[4][1] + K02[3] * A[4][2] + K03[3] * A[4][3]) * h, - ( - K[0][4] * A[4][0] - + K[1][4] * A[4][1] - + K[2][4] * A[4][2] - + K[3][4] * A[4][3] - ) + (K00[4] * A[4][0] + K01[4] * A[4][1] + K02[4] * A[4][2] + K03[4] * A[4][3]) * h, - ( - K[0][5] * A[4][0] - + K[1][5] * A[4][1] - + K[2][5] * A[4][2] - + K[3][5] * A[4][3] - ) + (K00[5] * A[4][0] + K01[5] * A[4][1] + K02[5] * A[4][2] + K03[5] * A[4][3]) * h, ] ) - K[4] = fun(t + C[4] * h, y + dy, argk) + K04 = fun(t + C[4] * h, y + dy, argk) # dy = np.dot(K[:5].T, A[5, :5]) * h dy = np.array( [ ( - K[0][0] * A[5][0] - + K[1][0] * A[5][1] - + K[2][0] * A[5][2] - + K[3][0] * A[5][3] - + K[4][0] * A[5][4] + K00[0] * A[5][0] + + K01[0] * A[5][1] + + K02[0] * A[5][2] + + K03[0] * A[5][3] + + K04[0] * A[5][4] ) * h, ( - K[0][1] * A[5][0] - + K[1][1] * A[5][1] - + K[2][1] * A[5][2] - + K[3][1] * A[5][3] - + K[4][1] * A[5][4] + K00[1] * A[5][0] + + K01[1] * A[5][1] + + K02[1] * A[5][2] + + K03[1] * A[5][3] + + K04[1] * A[5][4] ) * h, ( - K[0][2] * A[5][0] - + K[1][2] * A[5][1] - + K[2][2] * A[5][2] - + K[3][2] * A[5][3] - + K[4][2] * A[5][4] + K00[2] * A[5][0] + + K01[2] * A[5][1] + + K02[2] * A[5][2] + + K03[2] * A[5][3] + + K04[2] * A[5][4] ) * h, ( - K[0][3] * A[5][0] - + K[1][3] * A[5][1] - + K[2][3] * A[5][2] - + K[3][3] * A[5][3] - + K[4][3] * A[5][4] + K00[3] * A[5][0] + + K01[3] * A[5][1] + + K02[3] * A[5][2] + + K03[3] * A[5][3] + + K04[3] * A[5][4] ) * h, ( - K[0][4] * A[5][0] - + K[1][4] * A[5][1] - + K[2][4] * A[5][2] - + K[3][4] * A[5][3] - + K[4][4] * A[5][4] + K00[4] * A[5][0] + + K01[4] * A[5][1] + + K02[4] * A[5][2] + + K03[4] * A[5][3] + + K04[4] * A[5][4] ) * h, ( - K[0][5] * A[5][0] - + K[1][5] * A[5][1] - + K[2][5] * A[5][2] - + K[3][5] * A[5][3] - + K[4][5] * A[5][4] + K00[5] * A[5][0] + + K01[5] * A[5][1] + + K02[5] * A[5][2] + + K03[5] * A[5][3] + + K04[5] * A[5][4] ) * h, ] ) - K[5] = fun(t + C[5] * h, y + dy, argk) + K05 = fun(t + C[5] * h, y + dy, argk) # dy = np.dot(K[:6].T, A[6, :6]) * h dy = np.array( [ ( - K[0][0] * A[6][0] - + K[1][0] * A[6][1] - + K[2][0] * A[6][2] - + K[3][0] * A[6][3] - + K[4][0] * A[6][4] - + K[5][0] * A[6][5] + K00[0] * A[6][0] + + K01[0] * A[6][1] + + K02[0] * A[6][2] + + K03[0] * A[6][3] + + K04[0] * A[6][4] + + K05[0] * A[6][5] ) * h, ( - K[0][1] * A[6][0] - + K[1][1] * A[6][1] - + K[2][1] * A[6][2] - + K[3][1] * A[6][3] - + K[4][1] * A[6][4] - + K[5][1] * A[6][5] + K00[1] * A[6][0] + + K01[1] * A[6][1] + + K02[1] * A[6][2] + + K03[1] * A[6][3] + + K04[1] * A[6][4] + + K05[1] * A[6][5] ) * h, ( - K[0][2] * A[6][0] - + K[1][2] * A[6][1] - + K[2][2] * A[6][2] - + K[3][2] * A[6][3] - + K[4][2] * A[6][4] - + K[5][2] * A[6][5] + K00[2] * A[6][0] + + K01[2] * A[6][1] + + K02[2] * A[6][2] + + K03[2] * A[6][3] + + K04[2] * A[6][4] + + K05[2] * A[6][5] ) * h, ( - K[0][3] * A[6][0] - + K[1][3] * A[6][1] - + K[2][3] * A[6][2] - + K[3][3] * A[6][3] - + K[4][3] * A[6][4] - + K[5][3] * A[6][5] + K00[3] * A[6][0] + + K01[3] * A[6][1] + + K02[3] * A[6][2] + + K03[3] * A[6][3] + + K04[3] * A[6][4] + + K05[3] * A[6][5] ) * h, ( - K[0][4] * A[6][0] - + K[1][4] * A[6][1] - + K[2][4] * A[6][2] - + K[3][4] * A[6][3] - + K[4][4] * A[6][4] - + K[5][4] * A[6][5] + K00[4] * A[6][0] + + K01[4] * A[6][1] + + K02[4] * A[6][2] + + K03[4] * A[6][3] + + K04[4] * A[6][4] + + K05[4] * A[6][5] ) * h, ( - K[0][5] * A[6][0] - + K[1][5] * A[6][1] - + K[2][5] * A[6][2] - + K[3][5] * A[6][3] - + K[4][5] * A[6][4] - + K[5][5] * A[6][5] + K00[5] * A[6][0] + + K01[5] * A[6][1] + + K02[5] * A[6][2] + + K03[5] * A[6][3] + + K04[5] * A[6][4] + + K05[5] * A[6][5] ) * h, ] ) - K[6] = fun(t + C[6] * h, y + dy, argk) + K06 = fun(t + C[6] * h, y + dy, argk) # dy = np.dot(K[:7].T, A[7, :7]) * h dy = np.array( [ ( - K[0][0] * A[7][0] - + K[1][0] * A[7][1] - + K[2][0] * A[7][2] - + K[3][0] * A[7][3] - + K[4][0] * A[7][4] - + K[5][0] * A[7][5] - + K[6][0] * A[7][6] + K00[0] * A[7][0] + + K01[0] * A[7][1] + + K02[0] * A[7][2] + + K03[0] * A[7][3] + + K04[0] * A[7][4] + + K05[0] * A[7][5] + + K06[0] * A[7][6] ) * h, ( - K[0][1] * A[7][0] - + K[1][1] * A[7][1] - + K[2][1] * A[7][2] - + K[3][1] * A[7][3] - + K[4][1] * A[7][4] - + K[5][1] * A[7][5] - + K[6][1] * A[7][6] + K00[1] * A[7][0] + + K01[1] * A[7][1] + + K02[1] * A[7][2] + + K03[1] * A[7][3] + + K04[1] * A[7][4] + + K05[1] * A[7][5] + + K06[1] * A[7][6] ) * h, ( - K[0][2] * A[7][0] - + K[1][2] * A[7][1] - + K[2][2] * A[7][2] - + K[3][2] * A[7][3] - + K[4][2] * A[7][4] - + K[5][2] * A[7][5] - + K[6][2] * A[7][6] + K00[2] * A[7][0] + + K01[2] * A[7][1] + + K02[2] * A[7][2] + + K03[2] * A[7][3] + + K04[2] * A[7][4] + + K05[2] * A[7][5] + + K06[2] * A[7][6] ) * h, ( - K[0][3] * A[7][0] - + K[1][3] * A[7][1] - + K[2][3] * A[7][2] - + K[3][3] * A[7][3] - + K[4][3] * A[7][4] - + K[5][3] * A[7][5] - + K[6][3] * A[7][6] + K00[3] * A[7][0] + + K01[3] * A[7][1] + + K02[3] * A[7][2] + + K03[3] * A[7][3] + + K04[3] * A[7][4] + + K05[3] * A[7][5] + + K06[3] * A[7][6] ) * h, ( - K[0][4] * A[7][0] - + K[1][4] * A[7][1] - + K[2][4] * A[7][2] - + K[3][4] * A[7][3] - + K[4][4] * A[7][4] - + K[5][4] * A[7][5] - + K[6][4] * A[7][6] + K00[4] * A[7][0] + + K01[4] * A[7][1] + + K02[4] * A[7][2] + + K03[4] * A[7][3] + + K04[4] * A[7][4] + + K05[4] * A[7][5] + + K06[4] * A[7][6] ) * h, ( - K[0][5] * A[7][0] - + K[1][5] * A[7][1] - + K[2][5] * A[7][2] - + K[3][5] * A[7][3] - + K[4][5] * A[7][4] - + K[5][5] * A[7][5] - + K[6][5] * A[7][6] + K00[5] * A[7][0] + + K01[5] * A[7][1] + + K02[5] * A[7][2] + + K03[5] * A[7][3] + + K04[5] * A[7][4] + + K05[5] * A[7][5] + + K06[5] * A[7][6] ) * h, ] ) - K[7] = fun(t + C[7] * h, y + dy, argk) + K07 = fun(t + C[7] * h, y + dy, argk) # dy = np.dot(K[:8].T, A[8, :8]) * h dy = np.array( [ ( - K[0][0] * A[8][0] - + K[1][0] * A[8][1] - + K[2][0] * A[8][2] - + K[3][0] * A[8][3] - + K[4][0] * A[8][4] - + K[5][0] * A[8][5] - + K[6][0] * A[8][6] - + K[7][0] * A[8][7] + K00[0] * A[8][0] + + K01[0] * A[8][1] + + K02[0] * A[8][2] + + K03[0] * A[8][3] + + K04[0] * A[8][4] + + K05[0] * A[8][5] + + K06[0] * A[8][6] + + K07[0] * A[8][7] ) * h, ( - K[0][1] * A[8][0] - + K[1][1] * A[8][1] - + K[2][1] * A[8][2] - + K[3][1] * A[8][3] - + K[4][1] * A[8][4] - + K[5][1] * A[8][5] - + K[6][1] * A[8][6] - + K[7][1] * A[8][7] + K00[1] * A[8][0] + + K01[1] * A[8][1] + + K02[1] * A[8][2] + + K03[1] * A[8][3] + + K04[1] * A[8][4] + + K05[1] * A[8][5] + + K06[1] * A[8][6] + + K07[1] * A[8][7] ) * h, ( - K[0][2] * A[8][0] - + K[1][2] * A[8][1] - + K[2][2] * A[8][2] - + K[3][2] * A[8][3] - + K[4][2] * A[8][4] - + K[5][2] * A[8][5] - + K[6][2] * A[8][6] - + K[7][2] * A[8][7] + K00[2] * A[8][0] + + K01[2] * A[8][1] + + K02[2] * A[8][2] + + K03[2] * A[8][3] + + K04[2] * A[8][4] + + K05[2] * A[8][5] + + K06[2] * A[8][6] + + K07[2] * A[8][7] ) * h, ( - K[0][3] * A[8][0] - + K[1][3] * A[8][1] - + K[2][3] * A[8][2] - + K[3][3] * A[8][3] - + K[4][3] * A[8][4] - + K[5][3] * A[8][5] - + K[6][3] * A[8][6] - + K[7][3] * A[8][7] + K00[3] * A[8][0] + + K01[3] * A[8][1] + + K02[3] * A[8][2] + + K03[3] * A[8][3] + + K04[3] * A[8][4] + + K05[3] * A[8][5] + + K06[3] * A[8][6] + + K07[3] * A[8][7] ) * h, ( - K[0][4] * A[8][0] - + K[1][4] * A[8][1] - + K[2][4] * A[8][2] - + K[3][4] * A[8][3] - + K[4][4] * A[8][4] - + K[5][4] * A[8][5] - + K[6][4] * A[8][6] - + K[7][4] * A[8][7] + K00[4] * A[8][0] + + K01[4] * A[8][1] + + K02[4] * A[8][2] + + K03[4] * A[8][3] + + K04[4] * A[8][4] + + K05[4] * A[8][5] + + K06[4] * A[8][6] + + K07[4] * A[8][7] ) * h, ( - K[0][5] * A[8][0] - + K[1][5] * A[8][1] - + K[2][5] * A[8][2] - + K[3][5] * A[8][3] - + K[4][5] * A[8][4] - + K[5][5] * A[8][5] - + K[6][5] * A[8][6] - + K[7][5] * A[8][7] + K00[5] * A[8][0] + + K01[5] * A[8][1] + + K02[5] * A[8][2] + + K03[5] * A[8][3] + + K04[5] * A[8][4] + + K05[5] * A[8][5] + + K06[5] * A[8][6] + + K07[5] * A[8][7] ) * h, ] ) - K[8] = fun(t + C[8] * h, y + dy, argk) + K08 = fun(t + C[8] * h, y + dy, argk) # dy = np.dot(K[:9].T, A[9, :9]) * h dy = np.array( [ ( - K[0][0] * A[9][0] - + K[1][0] * A[9][1] - + K[2][0] * A[9][2] - + K[3][0] * A[9][3] - + K[4][0] * A[9][4] - + K[5][0] * A[9][5] - + K[6][0] * A[9][6] - + K[7][0] * A[9][7] - + K[8][0] * A[9][8] + K00[0] * A[9][0] + + K01[0] * A[9][1] + + K02[0] * A[9][2] + + K03[0] * A[9][3] + + K04[0] * A[9][4] + + K05[0] * A[9][5] + + K06[0] * A[9][6] + + K07[0] * A[9][7] + + K08[0] * A[9][8] ) * h, ( - K[0][1] * A[9][0] - + K[1][1] * A[9][1] - + K[2][1] * A[9][2] - + K[3][1] * A[9][3] - + K[4][1] * A[9][4] - + K[5][1] * A[9][5] - + K[6][1] * A[9][6] - + K[7][1] * A[9][7] - + K[8][1] * A[9][8] + K00[1] * A[9][0] + + K01[1] * A[9][1] + + K02[1] * A[9][2] + + K03[1] * A[9][3] + + K04[1] * A[9][4] + + K05[1] * A[9][5] + + K06[1] * A[9][6] + + K07[1] * A[9][7] + + K08[1] * A[9][8] ) * h, ( - K[0][2] * A[9][0] - + K[1][2] * A[9][1] - + K[2][2] * A[9][2] - + K[3][2] * A[9][3] - + K[4][2] * A[9][4] - + K[5][2] * A[9][5] - + K[6][2] * A[9][6] - + K[7][2] * A[9][7] - + K[8][2] * A[9][8] + K00[2] * A[9][0] + + K01[2] * A[9][1] + + K02[2] * A[9][2] + + K03[2] * A[9][3] + + K04[2] * A[9][4] + + K05[2] * A[9][5] + + K06[2] * A[9][6] + + K07[2] * A[9][7] + + K08[2] * A[9][8] ) * h, ( - K[0][3] * A[9][0] - + K[1][3] * A[9][1] - + K[2][3] * A[9][2] - + K[3][3] * A[9][3] - + K[4][3] * A[9][4] - + K[5][3] * A[9][5] - + K[6][3] * A[9][6] - + K[7][3] * A[9][7] - + K[8][3] * A[9][8] + K00[3] * A[9][0] + + K01[3] * A[9][1] + + K02[3] * A[9][2] + + K03[3] * A[9][3] + + K04[3] * A[9][4] + + K05[3] * A[9][5] + + K06[3] * A[9][6] + + K07[3] * A[9][7] + + K08[3] * A[9][8] ) * h, ( - K[0][4] * A[9][0] - + K[1][4] * A[9][1] - + K[2][4] * A[9][2] - + K[3][4] * A[9][3] - + K[4][4] * A[9][4] - + K[5][4] * A[9][5] - + K[6][4] * A[9][6] - + K[7][4] * A[9][7] - + K[8][4] * A[9][8] + K00[4] * A[9][0] + + K01[4] * A[9][1] + + K02[4] * A[9][2] + + K03[4] * A[9][3] + + K04[4] * A[9][4] + + K05[4] * A[9][5] + + K06[4] * A[9][6] + + K07[4] * A[9][7] + + K08[4] * A[9][8] ) * h, ( - K[0][5] * A[9][0] - + K[1][5] * A[9][1] - + K[2][5] * A[9][2] - + K[3][5] * A[9][3] - + K[4][5] * A[9][4] - + K[5][5] * A[9][5] - + K[6][5] * A[9][6] - + K[7][5] * A[9][7] - + K[8][5] * A[9][8] + K00[5] * A[9][0] + + K01[5] * A[9][1] + + K02[5] * A[9][2] + + K03[5] * A[9][3] + + K04[5] * A[9][4] + + K05[5] * A[9][5] + + K06[5] * A[9][6] + + K07[5] * A[9][7] + + K08[5] * A[9][8] ) * h, ] ) - K[9] = fun(t + C[9] * h, y + dy, argk) + K09 = fun(t + C[9] * h, y + dy, argk) # dy = np.dot(K[:10].T, A[10, :10]) * h dy = np.array( [ ( - K[0][0] * A[10][0] - + K[1][0] * A[10][1] - + K[2][0] * A[10][2] - + K[3][0] * A[10][3] - + K[4][0] * A[10][4] - + K[5][0] * A[10][5] - + K[6][0] * A[10][6] - + K[7][0] * A[10][7] - + K[8][0] * A[10][8] - + K[9][0] * A[10][9] - ) - * h, - ( - K[0][1] * A[10][0] - + K[1][1] * A[10][1] - + K[2][1] * A[10][2] - + K[3][1] * A[10][3] - + K[4][1] * A[10][4] - + K[5][1] * A[10][5] - + K[6][1] * A[10][6] - + K[7][1] * A[10][7] - + K[8][1] * A[10][8] - + K[9][1] * A[10][9] - ) - * h, - ( - K[0][2] * A[10][0] - + K[1][2] * A[10][1] - + K[2][2] * A[10][2] - + K[3][2] * A[10][3] - + K[4][2] * A[10][4] - + K[5][2] * A[10][5] - + K[6][2] * A[10][6] - + K[7][2] * A[10][7] - + K[8][2] * A[10][8] - + K[9][2] * A[10][9] - ) - * h, - ( - K[0][3] * A[10][0] - + K[1][3] * A[10][1] - + K[2][3] * A[10][2] - + K[3][3] * A[10][3] - + K[4][3] * A[10][4] - + K[5][3] * A[10][5] - + K[6][3] * A[10][6] - + K[7][3] * A[10][7] - + K[8][3] * A[10][8] - + K[9][3] * A[10][9] - ) - * h, - ( - K[0][4] * A[10][0] - + K[1][4] * A[10][1] - + K[2][4] * A[10][2] - + K[3][4] * A[10][3] - + K[4][4] * A[10][4] - + K[5][4] * A[10][5] - + K[6][4] * A[10][6] - + K[7][4] * A[10][7] - + K[8][4] * A[10][8] - + K[9][4] * A[10][9] - ) - * h, - ( - K[0][5] * A[10][0] - + K[1][5] * A[10][1] - + K[2][5] * A[10][2] - + K[3][5] * A[10][3] - + K[4][5] * A[10][4] - + K[5][5] * A[10][5] - + K[6][5] * A[10][6] - + K[7][5] * A[10][7] - + K[8][5] * A[10][8] - + K[9][5] * A[10][9] + K00[0] * A[10][0] + + K01[0] * A[10][1] + + K02[0] * A[10][2] + + K03[0] * A[10][3] + + K04[0] * A[10][4] + + K05[0] * A[10][5] + + K06[0] * A[10][6] + + K07[0] * A[10][7] + + K08[0] * A[10][8] + + K09[0] * A[10][9] + ) + * h, + ( + K00[1] * A[10][0] + + K01[1] * A[10][1] + + K02[1] * A[10][2] + + K03[1] * A[10][3] + + K04[1] * A[10][4] + + K05[1] * A[10][5] + + K06[1] * A[10][6] + + K07[1] * A[10][7] + + K08[1] * A[10][8] + + K09[1] * A[10][9] + ) + * h, + ( + K00[2] * A[10][0] + + K01[2] * A[10][1] + + K02[2] * A[10][2] + + K03[2] * A[10][3] + + K04[2] * A[10][4] + + K05[2] * A[10][5] + + K06[2] * A[10][6] + + K07[2] * A[10][7] + + K08[2] * A[10][8] + + K09[2] * A[10][9] + ) + * h, + ( + K00[3] * A[10][0] + + K01[3] * A[10][1] + + K02[3] * A[10][2] + + K03[3] * A[10][3] + + K04[3] * A[10][4] + + K05[3] * A[10][5] + + K06[3] * A[10][6] + + K07[3] * A[10][7] + + K08[3] * A[10][8] + + K09[3] * A[10][9] + ) + * h, + ( + K00[4] * A[10][0] + + K01[4] * A[10][1] + + K02[4] * A[10][2] + + K03[4] * A[10][3] + + K04[4] * A[10][4] + + K05[4] * A[10][5] + + K06[4] * A[10][6] + + K07[4] * A[10][7] + + K08[4] * A[10][8] + + K09[4] * A[10][9] + ) + * h, + ( + K00[5] * A[10][0] + + K01[5] * A[10][1] + + K02[5] * A[10][2] + + K03[5] * A[10][3] + + K04[5] * A[10][4] + + K05[5] * A[10][5] + + K06[5] * A[10][6] + + K07[5] * A[10][7] + + K08[5] * A[10][8] + + K09[5] * A[10][9] ) * h, ] ) - K[10] = fun(t + C[10] * h, y + dy, argk) + K10 = fun(t + C[10] * h, y + dy, argk) # dy = np.dot(K[:11].T, A[11, :11]) * h dy = np.array( [ ( - K[0][0] * A[11][0] - + K[1][0] * A[11][1] - + K[2][0] * A[11][2] - + K[3][0] * A[11][3] - + K[4][0] * A[11][4] - + K[5][0] * A[11][5] - + K[6][0] * A[11][6] - + K[7][0] * A[11][7] - + K[8][0] * A[11][8] - + K[9][0] * A[11][9] - + K[10][0] * A[11][10] - ) - * h, - ( - K[0][1] * A[11][0] - + K[1][1] * A[11][1] - + K[2][1] * A[11][2] - + K[3][1] * A[11][3] - + K[4][1] * A[11][4] - + K[5][1] * A[11][5] - + K[6][1] * A[11][6] - + K[7][1] * A[11][7] - + K[8][1] * A[11][8] - + K[9][1] * A[11][9] - + K[10][1] * A[11][10] - ) - * h, - ( - K[0][2] * A[11][0] - + K[1][2] * A[11][1] - + K[2][2] * A[11][2] - + K[3][2] * A[11][3] - + K[4][2] * A[11][4] - + K[5][2] * A[11][5] - + K[6][2] * A[11][6] - + K[7][2] * A[11][7] - + K[8][2] * A[11][8] - + K[9][2] * A[11][9] - + K[10][2] * A[11][10] - ) - * h, - ( - K[0][3] * A[11][0] - + K[1][3] * A[11][1] - + K[2][3] * A[11][2] - + K[3][3] * A[11][3] - + K[4][3] * A[11][4] - + K[5][3] * A[11][5] - + K[6][3] * A[11][6] - + K[7][3] * A[11][7] - + K[8][3] * A[11][8] - + K[9][3] * A[11][9] - + K[10][3] * A[11][10] - ) - * h, - ( - K[0][4] * A[11][0] - + K[1][4] * A[11][1] - + K[2][4] * A[11][2] - + K[3][4] * A[11][3] - + K[4][4] * A[11][4] - + K[5][4] * A[11][5] - + K[6][4] * A[11][6] - + K[7][4] * A[11][7] - + K[8][4] * A[11][8] - + K[9][4] * A[11][9] - + K[10][4] * A[11][10] - ) - * h, - ( - K[0][5] * A[11][0] - + K[1][5] * A[11][1] - + K[2][5] * A[11][2] - + K[3][5] * A[11][3] - + K[4][5] * A[11][4] - + K[5][5] * A[11][5] - + K[6][5] * A[11][6] - + K[7][5] * A[11][7] - + K[8][5] * A[11][8] - + K[9][5] * A[11][9] - + K[10][5] * A[11][10] + K00[0] * A[11][0] + + K01[0] * A[11][1] + + K02[0] * A[11][2] + + K03[0] * A[11][3] + + K04[0] * A[11][4] + + K05[0] * A[11][5] + + K06[0] * A[11][6] + + K07[0] * A[11][7] + + K08[0] * A[11][8] + + K09[0] * A[11][9] + + K10[0] * A[11][10] + ) + * h, + ( + K00[1] * A[11][0] + + K01[1] * A[11][1] + + K02[1] * A[11][2] + + K03[1] * A[11][3] + + K04[1] * A[11][4] + + K05[1] * A[11][5] + + K06[1] * A[11][6] + + K07[1] * A[11][7] + + K08[1] * A[11][8] + + K09[1] * A[11][9] + + K10[1] * A[11][10] + ) + * h, + ( + K00[2] * A[11][0] + + K01[2] * A[11][1] + + K02[2] * A[11][2] + + K03[2] * A[11][3] + + K04[2] * A[11][4] + + K05[2] * A[11][5] + + K06[2] * A[11][6] + + K07[2] * A[11][7] + + K08[2] * A[11][8] + + K09[2] * A[11][9] + + K10[2] * A[11][10] + ) + * h, + ( + K00[3] * A[11][0] + + K01[3] * A[11][1] + + K02[3] * A[11][2] + + K03[3] * A[11][3] + + K04[3] * A[11][4] + + K05[3] * A[11][5] + + K06[3] * A[11][6] + + K07[3] * A[11][7] + + K08[3] * A[11][8] + + K09[3] * A[11][9] + + K10[3] * A[11][10] + ) + * h, + ( + K00[4] * A[11][0] + + K01[4] * A[11][1] + + K02[4] * A[11][2] + + K03[4] * A[11][3] + + K04[4] * A[11][4] + + K05[4] * A[11][5] + + K06[4] * A[11][6] + + K07[4] * A[11][7] + + K08[4] * A[11][8] + + K09[4] * A[11][9] + + K10[4] * A[11][10] + ) + * h, + ( + K00[5] * A[11][0] + + K01[5] * A[11][1] + + K02[5] * A[11][2] + + K03[5] * A[11][3] + + K04[5] * A[11][4] + + K05[5] * A[11][5] + + K06[5] * A[11][6] + + K07[5] * A[11][7] + + K08[5] * A[11][8] + + K09[5] * A[11][9] + + K10[5] * A[11][10] ) * h, ] ) - K[11] = fun(t + C[11] * h, y + dy, argk) + K11 = fun(t + C[11] * h, y + dy, argk) + + K_ = np.array( + [ + K00, + K01, + K02, + K03, + K04, + K05, + K06, + K07, + K08, + K09, + K10, + K11, + ] + ).T - y_new = y + h * np.dot(K[:-1].T, B) + y_new = y + h * np.dot(K_, B) f_new = fun(t + h, y_new, argk) - K[-1] = f_new + K12 = f_new assert y_new.shape == (N_RV,) assert f_new.shape == (N_RV,) - return y_new, f_new + return ( + y_new, + f_new, + ( + K00, + K01, + K02, + K03, + K04, + K05, + K06, + K07, + K08, + K09, + K10, + K11, + K12, + ), + ) From 16f28e55d0b425db687d5b98acb82a074f63b6ed Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 22 Jan 2024 12:44:26 +0100 Subject: [PATCH 147/346] assert shapes, swap dot to matmul --- src/hapsira/core/math/ivp/_rkstep.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index 4dafac978..da311a845 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -680,7 +680,10 @@ def rk_step( ] ).T - y_new = y + h * np.dot(K_, B) + assert K_.shape == (6, 12) + assert B.shape == (12,) + + y_new = y + h * K_ @ B # np.dot(K_, B) f_new = fun(t + h, y_new, argk) K12 = f_new From 59f022dfaac2552d6216588e377001b3b58f3b8d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 22 Jan 2024 13:07:35 +0100 Subject: [PATCH 148/346] rm matmul, swap for tuple ops --- src/hapsira/core/math/ivp/_rkstep.py | 131 +++++++++++++++++++++++---- 1 file changed, 113 insertions(+), 18 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index da311a845..eca6a5645 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -663,27 +663,122 @@ def rk_step( ) K11 = fun(t + C[11] * h, y + dy, argk) - K_ = np.array( + # K_ = np.array( + # [ + # K00, + # K01, + # K02, + # K03, + # K04, + # K05, + # K06, + # K07, + # K08, + # K09, + # K10, + # K11, + # ] + # ) + + # assert K_.shape == (12, 6) + # assert B.shape == (12,) + + dy = np.array( [ - K00, - K01, - K02, - K03, - K04, - K05, - K06, - K07, - K08, - K09, - K10, - K11, + ( + K00[0] * B[0] + + K01[0] * B[1] + + K02[0] * B[2] + + K03[0] * B[3] + + K04[0] * B[4] + + K05[0] * B[5] + + K06[0] * B[6] + + K07[0] * B[7] + + K08[0] * B[8] + + K09[0] * B[9] + + K10[0] * B[10] + + K11[0] * B[11] + ) + * h, + ( + K00[1] * B[0] + + K01[1] * B[1] + + K02[1] * B[2] + + K03[1] * B[3] + + K04[1] * B[4] + + K05[1] * B[5] + + K06[1] * B[6] + + K07[1] * B[7] + + K08[1] * B[8] + + K09[1] * B[9] + + K10[1] * B[10] + + K11[1] * B[11] + ) + * h, + ( + K00[2] * B[0] + + K01[2] * B[1] + + K02[2] * B[2] + + K03[2] * B[3] + + K04[2] * B[4] + + K05[2] * B[5] + + K06[2] * B[6] + + K07[2] * B[7] + + K08[2] * B[8] + + K09[2] * B[9] + + K10[2] * B[10] + + K11[2] * B[11] + ) + * h, + ( + K00[3] * B[0] + + K01[3] * B[1] + + K02[3] * B[2] + + K03[3] * B[3] + + K04[3] * B[4] + + K05[3] * B[5] + + K06[3] * B[6] + + K07[3] * B[7] + + K08[3] * B[8] + + K09[3] * B[9] + + K10[3] * B[10] + + K11[3] * B[11] + ) + * h, + ( + K00[4] * B[0] + + K01[4] * B[1] + + K02[4] * B[2] + + K03[4] * B[3] + + K04[4] * B[4] + + K05[4] * B[5] + + K06[4] * B[6] + + K07[4] * B[7] + + K08[4] * B[8] + + K09[4] * B[9] + + K10[4] * B[10] + + K11[4] * B[11] + ) + * h, + ( + K00[5] * B[0] + + K01[4] * B[1] + + K02[5] * B[2] + + K03[5] * B[3] + + K04[5] * B[4] + + K05[5] * B[5] + + K06[5] * B[6] + + K07[5] * B[7] + + K08[5] * B[8] + + K09[5] * B[9] + + K10[5] * B[10] + + K11[5] * B[11] + ) + * h, ] - ).T - - assert K_.shape == (6, 12) - assert B.shape == (12,) + ) - y_new = y + h * K_ @ B # np.dot(K_, B) + y_new = y + dy # h * K_ @ B f_new = fun(t + h, y_new, argk) K12 = f_new From 545745cb9f4e63c19c2fad4ae4168873b1a7c790 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 22 Jan 2024 13:10:50 +0100 Subject: [PATCH 149/346] rm comments --- src/hapsira/core/math/ivp/_rkstep.py | 38 +--------------------------- 1 file changed, 1 insertion(+), 37 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index eca6a5645..5143efa57 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -88,11 +88,6 @@ def rk_step( K00 = f - # for s, (a, c) in enumerate(zip(A[1:], C[1:]), start=1): - # dy = np.dot(K[:s].T, a[:s]) * h - # K[s] = fun(t + c * h, y + dy) - - # dy = np.dot(K[:1].T, A[1, :1]) * h dy = np.array( [ (K00[0] * A[1][0]) * h, @@ -105,7 +100,6 @@ def rk_step( ) K01 = fun(t + C[1] * h, y + dy, argk) - # dy = np.dot(K[:2].T, A[2, :2]) * h dy = np.array( [ (K00[0] * A[2][0] + K01[0] * A[2][1]) * h, @@ -118,7 +112,6 @@ def rk_step( ) K02 = fun(t + C[2] * h, y + dy, argk) - # dy = np.dot(K[:3].T, A[3, :3]) * h dy = np.array( [ (K00[0] * A[3][0] + K01[0] * A[3][1] + K02[0] * A[3][2]) * h, @@ -132,7 +125,6 @@ def rk_step( K03 = fun(t + C[3] * h, y + dy, argk) - # dy = np.dot(K[:4].T, A[4, :4]) * h dy = np.array( [ (K00[0] * A[4][0] + K01[0] * A[4][1] + K02[0] * A[4][2] + K03[0] * A[4][3]) @@ -151,7 +143,6 @@ def rk_step( ) K04 = fun(t + C[4] * h, y + dy, argk) - # dy = np.dot(K[:5].T, A[5, :5]) * h dy = np.array( [ ( @@ -206,7 +197,6 @@ def rk_step( ) K05 = fun(t + C[5] * h, y + dy, argk) - # dy = np.dot(K[:6].T, A[6, :6]) * h dy = np.array( [ ( @@ -267,7 +257,6 @@ def rk_step( ) K06 = fun(t + C[6] * h, y + dy, argk) - # dy = np.dot(K[:7].T, A[7, :7]) * h dy = np.array( [ ( @@ -335,7 +324,6 @@ def rk_step( K07 = fun(t + C[7] * h, y + dy, argk) - # dy = np.dot(K[:8].T, A[8, :8]) * h dy = np.array( [ ( @@ -408,7 +396,6 @@ def rk_step( ) K08 = fun(t + C[8] * h, y + dy, argk) - # dy = np.dot(K[:9].T, A[9, :9]) * h dy = np.array( [ ( @@ -487,7 +474,6 @@ def rk_step( ) K09 = fun(t + C[9] * h, y + dy, argk) - # dy = np.dot(K[:10].T, A[10, :10]) * h dy = np.array( [ ( @@ -572,7 +558,6 @@ def rk_step( ) K10 = fun(t + C[10] * h, y + dy, argk) - # dy = np.dot(K[:11].T, A[11, :11]) * h dy = np.array( [ ( @@ -663,26 +648,6 @@ def rk_step( ) K11 = fun(t + C[11] * h, y + dy, argk) - # K_ = np.array( - # [ - # K00, - # K01, - # K02, - # K03, - # K04, - # K05, - # K06, - # K07, - # K08, - # K09, - # K10, - # K11, - # ] - # ) - - # assert K_.shape == (12, 6) - # assert B.shape == (12,) - dy = np.array( [ ( @@ -777,8 +742,7 @@ def rk_step( * h, ] ) - - y_new = y + dy # h * K_ @ B + y_new = y + dy f_new = fun(t + h, y_new, argk) K12 = f_new From d6469e300e2f6d331a581508e8cadf1b63fca09d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 22 Jan 2024 16:58:50 +0100 Subject: [PATCH 150/346] attempt nopython for rkstep, failing --- src/hapsira/core/jit.py | 13 + src/hapsira/core/math/ivp/_rk.py | 42 +- src/hapsira/core/math/ivp/_rkstep.py | 1425 +++++++++++---------- src/hapsira/core/propagation/__init__.py | 35 - src/hapsira/core/propagation/base.py | 24 +- src/hapsira/core/propagation/cowell.py | 17 +- src/hapsira/earth/__init__.py | 59 +- src/hapsira/twobody/propagation/cowell.py | 6 +- tests/tests_twobody/test_events.py | 35 +- tests/tests_twobody/test_perturbations.py | 323 ++--- tests/tests_twobody/test_propagation.py | 25 +- tests/tests_twobody/test_thrust.py | 81 +- 12 files changed, 1089 insertions(+), 996 deletions(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 537969064..29a2da490 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -10,6 +10,7 @@ __all__ = [ "PRECISIONS", + "DSIG", "hjit", "vjit", "gjit", @@ -27,6 +28,8 @@ PRECISIONS = ("f4", "f8") # TODO allow f2, i.e. half, for CUDA at least? +DSIG = "Tuple([V,V])(f,V,V,f)" + def _parse_signatures(signature: str, noreturn: bool = False) -> Union[str, List[str]]: """ @@ -126,6 +129,16 @@ def wrapper(inner_func: Callable) -> Callable: return wrapper +def djit(func): + """ + Wrapper for hjit to track differential equations + """ + + compiled = hjit(DSIG)(func) + compiled.djit = None # for debugging + return compiled + + def vjit(*args, **kwargs) -> Callable: """ Vectorize on array, pre-configured, user-facing, switches compiler targets. diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 428b78042..ddbb9de5a 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -5,7 +5,9 @@ from numba import jit from . import _dop853_coefficients as dop853_coefficients -from ._rkstep import rk_step, N_RV, N_STAGES +from ._rkstep import rk_step_hf, N_RV, N_STAGES + +from ...jit import array_to_V_hf __all__ = [ "EPS", @@ -90,7 +92,13 @@ def select_initial_step( h0 = 0.01 * d0 / d1 y1 = y0 + h0 * direction * f0 - f1 = fun(t0 + h0 * direction, y1, argk) + rr, vv = fun( + t0 + h0 * direction, + array_to_V_hf(y1[:3]), + array_to_V_hf(y1[3:]), + argk, + ) # TODO call into hf + f1 = np.array([*rr, *vv]) d2 = norm((f1 - f0) / scale) / h0 if d1 <= 1e-15 and d2 <= 1e-15: @@ -306,7 +314,13 @@ def __init__( self.y_old = None self.max_step = validate_max_step(max_step) self.rtol, self.atol = validate_tol(rtol, atol) - self.f = self.fun(self.t, self.y, self.argk) + rr, vv = self.fun( + self.t, + array_to_V_hf(self.y[:3]), + array_to_V_hf(self.y[3:]), + self.argk, + ) # TODO call into hf + self.f = np.array([*rr, *vv]) self.h_abs = select_initial_step( self.fun, self.t, @@ -420,7 +434,18 @@ def _step_impl(self): h = t_new - t h_abs = np.abs(h) - y_new, f_new, K_new = rk_step(self.fun, t, y, self.f, h, self.argk) + rr_new, vv_new, fr_new, fv_new, K_new = rk_step_hf( + self.fun, + t, + array_to_V_hf(y[:3]), + array_to_V_hf(y[3:]), + array_to_V_hf(self.f[:3]), + array_to_V_hf(self.f[3:]), + h, + self.argk, + ) # TODO call into hf + y_new = np.array([*rr_new, *vv_new]) + f_new = np.array([*fr_new, *fv_new]) self.K[: N_STAGES + 1, :N_RV] = np.array([K_new]) scale = atol + np.maximum(np.abs(y), np.abs(y_new)) * rtol @@ -458,7 +483,14 @@ def _dense_output_impl(self): h = self.h_previous for s, (a, c) in enumerate(zip(self.A_EXTRA, self.C_EXTRA), start=N_STAGES + 1): dy = np.dot(K[:s].T, a[:s]) * h - K[s] = self.fun(self.t_old + c * h, self.y_old + dy, self.argk) + y_ = self.y_old + dy + rr, vv = self.fun( + self.t_old + c * h, + array_to_V_hf(y_[:3]), + array_to_V_hf(y_[3:]), + self.argk, + ) # TODO call into hf + K[s] = np.array([*rr, *vv]) F = np.empty((INTERPOLATOR_POWER, N_RV), dtype=self.y_old.dtype) diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index 5143efa57..3ac459caa 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -1,12 +1,11 @@ from typing import Callable, Tuple -import numpy as np -from numba import jit - -from ._dop853_coefficients import A, B, C +from ._dop853_coefficients import A as _A, B as _B, C as _C +from ...math.linalg import add_VV_hf +from ...jit import hjit, DSIG __all__ = [ - "rk_step", + "rk_step_hf", "N_RV", "N_STAGES", ] @@ -14,20 +13,39 @@ N_RV = 6 N_STAGES = 12 -A = tuple(tuple(line) for line in A[:N_STAGES, :N_STAGES]) -# B = B -C = tuple(C[:N_STAGES]) - - -@jit(nopython=False) -def rk_step( +# A = tuple(tuple(float(number) for number in line) for line in A[:N_STAGES, :N_STAGES]) +A01 = tuple(float(number) for number in _A[1, :N_STAGES]) +A02 = tuple(float(number) for number in _A[2, :N_STAGES]) +A03 = tuple(float(number) for number in _A[3, :N_STAGES]) +A04 = tuple(float(number) for number in _A[4, :N_STAGES]) +A05 = tuple(float(number) for number in _A[5, :N_STAGES]) +A06 = tuple(float(number) for number in _A[6, :N_STAGES]) +A07 = tuple(float(number) for number in _A[7, :N_STAGES]) +A08 = tuple(float(number) for number in _A[8, :N_STAGES]) +A09 = tuple(float(number) for number in _A[9, :N_STAGES]) +A10 = tuple(float(number) for number in _A[10, :N_STAGES]) +A11 = tuple(float(number) for number in _A[11, :N_STAGES]) +B = tuple(float(number) for number in _B) +C = tuple(float(number) for number in _C[:N_STAGES]) + +_KSIG = ( + "Tuple([" + + ",".join(["Tuple([" + ",".join(["f"] * N_RV) + "])"] * (N_STAGES + 1)) + + "])" +) + + +@hjit(f"Tuple([V,V,V,V,{_KSIG:s}])(F({DSIG:s}),f,V,V,V,V,f,f)") +def rk_step_hf( fun: Callable, t: float, - y: np.ndarray, - f: np.ndarray, + rr: tuple[float, float, float], + vv: tuple[float, float, float], + fr: tuple[float, float, float], + fv: tuple[float, float, float], h: float, argk: float, -) -> Tuple[np.ndarray, np.ndarray]: +) -> Tuple[Tuple, Tuple, Tuple, Tuple, Tuple]: """Perform a single Runge-Kutta step. This function computes a prediction of an explicit Runge-Kutta method and @@ -41,9 +59,13 @@ def rk_step( Right-hand side of the system. t : float Current time. - y : ndarray, shape (n,) - Current state. - f : ndarray, shape (n,) + r : tuple[float,float,float] + Current r. + v : tuple[float,float,float] + Current v. + fr : tuple[float,float,float] + Current value of the derivative, i.e., ``fun(x, y)``. + fv : tuple[float,float,float] Current value of the derivative, i.e., ``fun(x, y)``. h : float Step to use. @@ -78,681 +100,730 @@ def rk_step( Equations I: Nonstiff Problems", Sec. II.4. """ - assert y.shape == (N_RV,) - assert f.shape == (N_RV,) - # assert K.shape == (N_STAGES + 1, N_RV) - - # assert A.shape == (N_STAGES, N_STAGES) - assert B.shape == (N_STAGES,) - # assert C.shape == (N_STAGES,) - - K00 = f + K00 = *fr, *fv - dy = np.array( - [ - (K00[0] * A[1][0]) * h, - (K00[1] * A[1][0]) * h, - (K00[2] * A[1][0]) * h, - (K00[3] * A[1][0]) * h, - (K00[4] * A[1][0]) * h, - (K00[5] * A[1][0]) * h, - ] - ) - K01 = fun(t + C[1] * h, y + dy, argk) - - dy = np.array( - [ - (K00[0] * A[2][0] + K01[0] * A[2][1]) * h, - (K00[1] * A[2][0] + K01[1] * A[2][1]) * h, - (K00[2] * A[2][0] + K01[2] * A[2][1]) * h, - (K00[3] * A[2][0] + K01[3] * A[2][1]) * h, - (K00[4] * A[2][0] + K01[4] * A[2][1]) * h, - (K00[5] * A[2][0] + K01[5] * A[2][1]) * h, - ] - ) - K02 = fun(t + C[2] * h, y + dy, argk) - - dy = np.array( - [ - (K00[0] * A[3][0] + K01[0] * A[3][1] + K02[0] * A[3][2]) * h, - (K00[1] * A[3][0] + K01[1] * A[3][1] + K02[1] * A[3][2]) * h, - (K00[2] * A[3][0] + K01[2] * A[3][1] + K02[2] * A[3][2]) * h, - (K00[3] * A[3][0] + K01[3] * A[3][1] + K02[3] * A[3][2]) * h, - (K00[4] * A[3][0] + K01[4] * A[3][1] + K02[4] * A[3][2]) * h, - (K00[5] * A[3][0] + K01[5] * A[3][1] + K02[5] * A[3][2]) * h, - ] + dr = ( + (K00[0] * A01[0]) * h, + (K00[1] * A01[0]) * h, + (K00[2] * A01[0]) * h, ) + dv = ( + (K00[3] * A01[0]) * h, + (K00[4] * A01[0]) * h, + (K00[5] * A01[0]) * h, + ) + fr, fv = fun( + t + C[1] * h, + add_VV_hf(rr, dr), + add_VV_hf(vv, dv), + argk, + ) + K01 = *fr, *fv - K03 = fun(t + C[3] * h, y + dy, argk) - - dy = np.array( - [ - (K00[0] * A[4][0] + K01[0] * A[4][1] + K02[0] * A[4][2] + K03[0] * A[4][3]) - * h, - (K00[1] * A[4][0] + K01[1] * A[4][1] + K02[1] * A[4][2] + K03[1] * A[4][3]) - * h, - (K00[2] * A[4][0] + K01[2] * A[4][1] + K02[2] * A[4][2] + K03[2] * A[4][3]) - * h, - (K00[3] * A[4][0] + K01[3] * A[4][1] + K02[3] * A[4][2] + K03[3] * A[4][3]) - * h, - (K00[4] * A[4][0] + K01[4] * A[4][1] + K02[4] * A[4][2] + K03[4] * A[4][3]) - * h, - (K00[5] * A[4][0] + K01[5] * A[4][1] + K02[5] * A[4][2] + K03[5] * A[4][3]) - * h, - ] - ) - K04 = fun(t + C[4] * h, y + dy, argk) - - dy = np.array( - [ - ( - K00[0] * A[5][0] - + K01[0] * A[5][1] - + K02[0] * A[5][2] - + K03[0] * A[5][3] - + K04[0] * A[5][4] - ) - * h, - ( - K00[1] * A[5][0] - + K01[1] * A[5][1] - + K02[1] * A[5][2] - + K03[1] * A[5][3] - + K04[1] * A[5][4] - ) - * h, - ( - K00[2] * A[5][0] - + K01[2] * A[5][1] - + K02[2] * A[5][2] - + K03[2] * A[5][3] - + K04[2] * A[5][4] - ) - * h, - ( - K00[3] * A[5][0] - + K01[3] * A[5][1] - + K02[3] * A[5][2] - + K03[3] * A[5][3] - + K04[3] * A[5][4] - ) - * h, - ( - K00[4] * A[5][0] - + K01[4] * A[5][1] - + K02[4] * A[5][2] - + K03[4] * A[5][3] - + K04[4] * A[5][4] - ) - * h, - ( - K00[5] * A[5][0] - + K01[5] * A[5][1] - + K02[5] * A[5][2] - + K03[5] * A[5][3] - + K04[5] * A[5][4] - ) - * h, - ] - ) - K05 = fun(t + C[5] * h, y + dy, argk) + dr = ( + (K00[0] * A02[0] + K01[0] * A02[1]) * h, + (K00[1] * A02[0] + K01[1] * A02[1]) * h, + (K00[2] * A02[0] + K01[2] * A02[1]) * h, + ) + dv = ( + (K00[3] * A02[0] + K01[3] * A02[1]) * h, + (K00[4] * A02[0] + K01[4] * A02[1]) * h, + (K00[5] * A02[0] + K01[5] * A02[1]) * h, + ) + fr, fv = fun( + t + C[2] * h, + add_VV_hf(rr, dr), + add_VV_hf(vv, dv), + argk, + ) + K02 = *fr, *fv - dy = np.array( - [ - ( - K00[0] * A[6][0] - + K01[0] * A[6][1] - + K02[0] * A[6][2] - + K03[0] * A[6][3] - + K04[0] * A[6][4] - + K05[0] * A[6][5] - ) - * h, - ( - K00[1] * A[6][0] - + K01[1] * A[6][1] - + K02[1] * A[6][2] - + K03[1] * A[6][3] - + K04[1] * A[6][4] - + K05[1] * A[6][5] - ) - * h, - ( - K00[2] * A[6][0] - + K01[2] * A[6][1] - + K02[2] * A[6][2] - + K03[2] * A[6][3] - + K04[2] * A[6][4] - + K05[2] * A[6][5] - ) - * h, - ( - K00[3] * A[6][0] - + K01[3] * A[6][1] - + K02[3] * A[6][2] - + K03[3] * A[6][3] - + K04[3] * A[6][4] - + K05[3] * A[6][5] - ) - * h, - ( - K00[4] * A[6][0] - + K01[4] * A[6][1] - + K02[4] * A[6][2] - + K03[4] * A[6][3] - + K04[4] * A[6][4] - + K05[4] * A[6][5] - ) - * h, - ( - K00[5] * A[6][0] - + K01[5] * A[6][1] - + K02[5] * A[6][2] - + K03[5] * A[6][3] - + K04[5] * A[6][4] - + K05[5] * A[6][5] - ) - * h, - ] - ) - K06 = fun(t + C[6] * h, y + dy, argk) + dr = ( + (K00[0] * A03[0] + K01[0] * A03[1] + K02[0] * A03[2]) * h, + (K00[1] * A03[0] + K01[1] * A03[1] + K02[1] * A03[2]) * h, + (K00[2] * A03[0] + K01[2] * A03[1] + K02[2] * A03[2]) * h, + ) + dv = ( + (K00[3] * A03[0] + K01[3] * A03[1] + K02[3] * A03[2]) * h, + (K00[4] * A03[0] + K01[4] * A03[1] + K02[4] * A03[2]) * h, + (K00[5] * A03[0] + K01[5] * A03[1] + K02[5] * A03[2]) * h, + ) + fr, fv = fun( + t + C[3] * h, + add_VV_hf(rr, dr), + add_VV_hf(vv, dv), + argk, + ) + K03 = *fr, *fv - dy = np.array( - [ - ( - K00[0] * A[7][0] - + K01[0] * A[7][1] - + K02[0] * A[7][2] - + K03[0] * A[7][3] - + K04[0] * A[7][4] - + K05[0] * A[7][5] - + K06[0] * A[7][6] - ) - * h, - ( - K00[1] * A[7][0] - + K01[1] * A[7][1] - + K02[1] * A[7][2] - + K03[1] * A[7][3] - + K04[1] * A[7][4] - + K05[1] * A[7][5] - + K06[1] * A[7][6] - ) - * h, - ( - K00[2] * A[7][0] - + K01[2] * A[7][1] - + K02[2] * A[7][2] - + K03[2] * A[7][3] - + K04[2] * A[7][4] - + K05[2] * A[7][5] - + K06[2] * A[7][6] - ) - * h, - ( - K00[3] * A[7][0] - + K01[3] * A[7][1] - + K02[3] * A[7][2] - + K03[3] * A[7][3] - + K04[3] * A[7][4] - + K05[3] * A[7][5] - + K06[3] * A[7][6] - ) - * h, - ( - K00[4] * A[7][0] - + K01[4] * A[7][1] - + K02[4] * A[7][2] - + K03[4] * A[7][3] - + K04[4] * A[7][4] - + K05[4] * A[7][5] - + K06[4] * A[7][6] - ) - * h, - ( - K00[5] * A[7][0] - + K01[5] * A[7][1] - + K02[5] * A[7][2] - + K03[5] * A[7][3] - + K04[5] * A[7][4] - + K05[5] * A[7][5] - + K06[5] * A[7][6] - ) - * h, - ] + dr = ( + (K00[0] * A04[0] + K01[0] * A04[1] + K02[0] * A04[2] + K03[0] * A04[3]) * h, + (K00[1] * A04[0] + K01[1] * A04[1] + K02[1] * A04[2] + K03[1] * A04[3]) * h, + (K00[2] * A04[0] + K01[2] * A04[1] + K02[2] * A04[2] + K03[2] * A04[3]) * h, ) + dv = ( + (K00[3] * A04[0] + K01[3] * A04[1] + K02[3] * A04[2] + K03[3] * A04[3]) * h, + (K00[4] * A04[0] + K01[4] * A04[1] + K02[4] * A04[2] + K03[4] * A04[3]) * h, + (K00[5] * A04[0] + K01[5] * A04[1] + K02[5] * A04[2] + K03[5] * A04[3]) * h, + ) + fr, fv = fun( + t + C[4] * h, + add_VV_hf(rr, dr), + add_VV_hf(vv, dv), + argk, + ) + K04 = *fr, *fv - K07 = fun(t + C[7] * h, y + dy, argk) + dr = ( + ( + K00[0] * A05[0] + + K01[0] * A05[1] + + K02[0] * A05[2] + + K03[0] * A05[3] + + K04[0] * A05[4] + ) + * h, + ( + K00[1] * A05[0] + + K01[1] * A05[1] + + K02[1] * A05[2] + + K03[1] * A05[3] + + K04[1] * A05[4] + ) + * h, + ( + K00[2] * A05[0] + + K01[2] * A05[1] + + K02[2] * A05[2] + + K03[2] * A05[3] + + K04[2] * A05[4] + ) + * h, + ) + dv = ( + ( + K00[3] * A05[0] + + K01[3] * A05[1] + + K02[3] * A05[2] + + K03[3] * A05[3] + + K04[3] * A05[4] + ) + * h, + ( + K00[4] * A05[0] + + K01[4] * A05[1] + + K02[4] * A05[2] + + K03[4] * A05[3] + + K04[4] * A05[4] + ) + * h, + ( + K00[5] * A05[0] + + K01[5] * A05[1] + + K02[5] * A05[2] + + K03[5] * A05[3] + + K04[5] * A05[4] + ) + * h, + ) + fr, fv = fun( + t + C[5] * h, + add_VV_hf(rr, dr), + add_VV_hf(vv, dv), + argk, + ) + K05 = *fr, *fv - dy = np.array( - [ - ( - K00[0] * A[8][0] - + K01[0] * A[8][1] - + K02[0] * A[8][2] - + K03[0] * A[8][3] - + K04[0] * A[8][4] - + K05[0] * A[8][5] - + K06[0] * A[8][6] - + K07[0] * A[8][7] - ) - * h, - ( - K00[1] * A[8][0] - + K01[1] * A[8][1] - + K02[1] * A[8][2] - + K03[1] * A[8][3] - + K04[1] * A[8][4] - + K05[1] * A[8][5] - + K06[1] * A[8][6] - + K07[1] * A[8][7] - ) - * h, - ( - K00[2] * A[8][0] - + K01[2] * A[8][1] - + K02[2] * A[8][2] - + K03[2] * A[8][3] - + K04[2] * A[8][4] - + K05[2] * A[8][5] - + K06[2] * A[8][6] - + K07[2] * A[8][7] - ) - * h, - ( - K00[3] * A[8][0] - + K01[3] * A[8][1] - + K02[3] * A[8][2] - + K03[3] * A[8][3] - + K04[3] * A[8][4] - + K05[3] * A[8][5] - + K06[3] * A[8][6] - + K07[3] * A[8][7] - ) - * h, - ( - K00[4] * A[8][0] - + K01[4] * A[8][1] - + K02[4] * A[8][2] - + K03[4] * A[8][3] - + K04[4] * A[8][4] - + K05[4] * A[8][5] - + K06[4] * A[8][6] - + K07[4] * A[8][7] - ) - * h, - ( - K00[5] * A[8][0] - + K01[5] * A[8][1] - + K02[5] * A[8][2] - + K03[5] * A[8][3] - + K04[5] * A[8][4] - + K05[5] * A[8][5] - + K06[5] * A[8][6] - + K07[5] * A[8][7] - ) - * h, - ] - ) - K08 = fun(t + C[8] * h, y + dy, argk) + dr = ( + ( + K00[0] * A06[0] + + K01[0] * A06[1] + + K02[0] * A06[2] + + K03[0] * A06[3] + + K04[0] * A06[4] + + K05[0] * A06[5] + ) + * h, + ( + K00[1] * A06[0] + + K01[1] * A06[1] + + K02[1] * A06[2] + + K03[1] * A06[3] + + K04[1] * A06[4] + + K05[1] * A06[5] + ) + * h, + ( + K00[2] * A06[0] + + K01[2] * A06[1] + + K02[2] * A06[2] + + K03[2] * A06[3] + + K04[2] * A06[4] + + K05[2] * A06[5] + ) + * h, + ) + dv = ( + ( + K00[3] * A06[0] + + K01[3] * A06[1] + + K02[3] * A06[2] + + K03[3] * A06[3] + + K04[3] * A06[4] + + K05[3] * A06[5] + ) + * h, + ( + K00[4] * A06[0] + + K01[4] * A06[1] + + K02[4] * A06[2] + + K03[4] * A06[3] + + K04[4] * A06[4] + + K05[4] * A06[5] + ) + * h, + ( + K00[5] * A06[0] + + K01[5] * A06[1] + + K02[5] * A06[2] + + K03[5] * A06[3] + + K04[5] * A06[4] + + K05[5] * A06[5] + ) + * h, + ) + fr, fv = fun( + t + C[6] * h, + add_VV_hf(rr, dr), + add_VV_hf(vv, dv), + argk, + ) + K06 = *fr, *fv - dy = np.array( - [ - ( - K00[0] * A[9][0] - + K01[0] * A[9][1] - + K02[0] * A[9][2] - + K03[0] * A[9][3] - + K04[0] * A[9][4] - + K05[0] * A[9][5] - + K06[0] * A[9][6] - + K07[0] * A[9][7] - + K08[0] * A[9][8] - ) - * h, - ( - K00[1] * A[9][0] - + K01[1] * A[9][1] - + K02[1] * A[9][2] - + K03[1] * A[9][3] - + K04[1] * A[9][4] - + K05[1] * A[9][5] - + K06[1] * A[9][6] - + K07[1] * A[9][7] - + K08[1] * A[9][8] - ) - * h, - ( - K00[2] * A[9][0] - + K01[2] * A[9][1] - + K02[2] * A[9][2] - + K03[2] * A[9][3] - + K04[2] * A[9][4] - + K05[2] * A[9][5] - + K06[2] * A[9][6] - + K07[2] * A[9][7] - + K08[2] * A[9][8] - ) - * h, - ( - K00[3] * A[9][0] - + K01[3] * A[9][1] - + K02[3] * A[9][2] - + K03[3] * A[9][3] - + K04[3] * A[9][4] - + K05[3] * A[9][5] - + K06[3] * A[9][6] - + K07[3] * A[9][7] - + K08[3] * A[9][8] - ) - * h, - ( - K00[4] * A[9][0] - + K01[4] * A[9][1] - + K02[4] * A[9][2] - + K03[4] * A[9][3] - + K04[4] * A[9][4] - + K05[4] * A[9][5] - + K06[4] * A[9][6] - + K07[4] * A[9][7] - + K08[4] * A[9][8] - ) - * h, - ( - K00[5] * A[9][0] - + K01[5] * A[9][1] - + K02[5] * A[9][2] - + K03[5] * A[9][3] - + K04[5] * A[9][4] - + K05[5] * A[9][5] - + K06[5] * A[9][6] - + K07[5] * A[9][7] - + K08[5] * A[9][8] - ) - * h, - ] - ) - K09 = fun(t + C[9] * h, y + dy, argk) + dr = ( + ( + K00[0] * A07[0] + + K01[0] * A07[1] + + K02[0] * A07[2] + + K03[0] * A07[3] + + K04[0] * A07[4] + + K05[0] * A07[5] + + K06[0] * A07[6] + ) + * h, + ( + K00[1] * A07[0] + + K01[1] * A07[1] + + K02[1] * A07[2] + + K03[1] * A07[3] + + K04[1] * A07[4] + + K05[1] * A07[5] + + K06[1] * A07[6] + ) + * h, + ( + K00[2] * A07[0] + + K01[2] * A07[1] + + K02[2] * A07[2] + + K03[2] * A07[3] + + K04[2] * A07[4] + + K05[2] * A07[5] + + K06[2] * A07[6] + ) + * h, + ) + dv = ( + ( + K00[3] * A07[0] + + K01[3] * A07[1] + + K02[3] * A07[2] + + K03[3] * A07[3] + + K04[3] * A07[4] + + K05[3] * A07[5] + + K06[3] * A07[6] + ) + * h, + ( + K00[4] * A07[0] + + K01[4] * A07[1] + + K02[4] * A07[2] + + K03[4] * A07[3] + + K04[4] * A07[4] + + K05[4] * A07[5] + + K06[4] * A07[6] + ) + * h, + ( + K00[5] * A07[0] + + K01[5] * A07[1] + + K02[5] * A07[2] + + K03[5] * A07[3] + + K04[5] * A07[4] + + K05[5] * A07[5] + + K06[5] * A07[6] + ) + * h, + ) + fr, fv = fun( + t + C[7] * h, + add_VV_hf(rr, dr), + add_VV_hf(vv, dv), + argk, + ) + K07 = *fr, *fv - dy = np.array( - [ - ( - K00[0] * A[10][0] - + K01[0] * A[10][1] - + K02[0] * A[10][2] - + K03[0] * A[10][3] - + K04[0] * A[10][4] - + K05[0] * A[10][5] - + K06[0] * A[10][6] - + K07[0] * A[10][7] - + K08[0] * A[10][8] - + K09[0] * A[10][9] - ) - * h, - ( - K00[1] * A[10][0] - + K01[1] * A[10][1] - + K02[1] * A[10][2] - + K03[1] * A[10][3] - + K04[1] * A[10][4] - + K05[1] * A[10][5] - + K06[1] * A[10][6] - + K07[1] * A[10][7] - + K08[1] * A[10][8] - + K09[1] * A[10][9] - ) - * h, - ( - K00[2] * A[10][0] - + K01[2] * A[10][1] - + K02[2] * A[10][2] - + K03[2] * A[10][3] - + K04[2] * A[10][4] - + K05[2] * A[10][5] - + K06[2] * A[10][6] - + K07[2] * A[10][7] - + K08[2] * A[10][8] - + K09[2] * A[10][9] - ) - * h, - ( - K00[3] * A[10][0] - + K01[3] * A[10][1] - + K02[3] * A[10][2] - + K03[3] * A[10][3] - + K04[3] * A[10][4] - + K05[3] * A[10][5] - + K06[3] * A[10][6] - + K07[3] * A[10][7] - + K08[3] * A[10][8] - + K09[3] * A[10][9] - ) - * h, - ( - K00[4] * A[10][0] - + K01[4] * A[10][1] - + K02[4] * A[10][2] - + K03[4] * A[10][3] - + K04[4] * A[10][4] - + K05[4] * A[10][5] - + K06[4] * A[10][6] - + K07[4] * A[10][7] - + K08[4] * A[10][8] - + K09[4] * A[10][9] - ) - * h, - ( - K00[5] * A[10][0] - + K01[5] * A[10][1] - + K02[5] * A[10][2] - + K03[5] * A[10][3] - + K04[5] * A[10][4] - + K05[5] * A[10][5] - + K06[5] * A[10][6] - + K07[5] * A[10][7] - + K08[5] * A[10][8] - + K09[5] * A[10][9] - ) - * h, - ] - ) - K10 = fun(t + C[10] * h, y + dy, argk) + dr = ( + ( + K00[0] * A08[0] + + K01[0] * A08[1] + + K02[0] * A08[2] + + K03[0] * A08[3] + + K04[0] * A08[4] + + K05[0] * A08[5] + + K06[0] * A08[6] + + K07[0] * A08[7] + ) + * h, + ( + K00[1] * A08[0] + + K01[1] * A08[1] + + K02[1] * A08[2] + + K03[1] * A08[3] + + K04[1] * A08[4] + + K05[1] * A08[5] + + K06[1] * A08[6] + + K07[1] * A08[7] + ) + * h, + ( + K00[2] * A08[0] + + K01[2] * A08[1] + + K02[2] * A08[2] + + K03[2] * A08[3] + + K04[2] * A08[4] + + K05[2] * A08[5] + + K06[2] * A08[6] + + K07[2] * A08[7] + ) + * h, + ) + dv = ( + ( + K00[3] * A08[0] + + K01[3] * A08[1] + + K02[3] * A08[2] + + K03[3] * A08[3] + + K04[3] * A08[4] + + K05[3] * A08[5] + + K06[3] * A08[6] + + K07[3] * A08[7] + ) + * h, + ( + K00[4] * A08[0] + + K01[4] * A08[1] + + K02[4] * A08[2] + + K03[4] * A08[3] + + K04[4] * A08[4] + + K05[4] * A08[5] + + K06[4] * A08[6] + + K07[4] * A08[7] + ) + * h, + ( + K00[5] * A08[0] + + K01[5] * A08[1] + + K02[5] * A08[2] + + K03[5] * A08[3] + + K04[5] * A08[4] + + K05[5] * A08[5] + + K06[5] * A08[6] + + K07[5] * A08[7] + ) + * h, + ) + fr, fv = fun( + t + C[8] * h, + add_VV_hf(rr, dr), + add_VV_hf(vv, dv), + argk, + ) + K08 = *fr, *fv - dy = np.array( - [ - ( - K00[0] * A[11][0] - + K01[0] * A[11][1] - + K02[0] * A[11][2] - + K03[0] * A[11][3] - + K04[0] * A[11][4] - + K05[0] * A[11][5] - + K06[0] * A[11][6] - + K07[0] * A[11][7] - + K08[0] * A[11][8] - + K09[0] * A[11][9] - + K10[0] * A[11][10] - ) - * h, - ( - K00[1] * A[11][0] - + K01[1] * A[11][1] - + K02[1] * A[11][2] - + K03[1] * A[11][3] - + K04[1] * A[11][4] - + K05[1] * A[11][5] - + K06[1] * A[11][6] - + K07[1] * A[11][7] - + K08[1] * A[11][8] - + K09[1] * A[11][9] - + K10[1] * A[11][10] - ) - * h, - ( - K00[2] * A[11][0] - + K01[2] * A[11][1] - + K02[2] * A[11][2] - + K03[2] * A[11][3] - + K04[2] * A[11][4] - + K05[2] * A[11][5] - + K06[2] * A[11][6] - + K07[2] * A[11][7] - + K08[2] * A[11][8] - + K09[2] * A[11][9] - + K10[2] * A[11][10] - ) - * h, - ( - K00[3] * A[11][0] - + K01[3] * A[11][1] - + K02[3] * A[11][2] - + K03[3] * A[11][3] - + K04[3] * A[11][4] - + K05[3] * A[11][5] - + K06[3] * A[11][6] - + K07[3] * A[11][7] - + K08[3] * A[11][8] - + K09[3] * A[11][9] - + K10[3] * A[11][10] - ) - * h, - ( - K00[4] * A[11][0] - + K01[4] * A[11][1] - + K02[4] * A[11][2] - + K03[4] * A[11][3] - + K04[4] * A[11][4] - + K05[4] * A[11][5] - + K06[4] * A[11][6] - + K07[4] * A[11][7] - + K08[4] * A[11][8] - + K09[4] * A[11][9] - + K10[4] * A[11][10] - ) - * h, - ( - K00[5] * A[11][0] - + K01[5] * A[11][1] - + K02[5] * A[11][2] - + K03[5] * A[11][3] - + K04[5] * A[11][4] - + K05[5] * A[11][5] - + K06[5] * A[11][6] - + K07[5] * A[11][7] - + K08[5] * A[11][8] - + K09[5] * A[11][9] - + K10[5] * A[11][10] - ) - * h, - ] - ) - K11 = fun(t + C[11] * h, y + dy, argk) + dr = ( + ( + K00[0] * A09[0] + + K01[0] * A09[1] + + K02[0] * A09[2] + + K03[0] * A09[3] + + K04[0] * A09[4] + + K05[0] * A09[5] + + K06[0] * A09[6] + + K07[0] * A09[7] + + K08[0] * A09[8] + ) + * h, + ( + K00[1] * A09[0] + + K01[1] * A09[1] + + K02[1] * A09[2] + + K03[1] * A09[3] + + K04[1] * A09[4] + + K05[1] * A09[5] + + K06[1] * A09[6] + + K07[1] * A09[7] + + K08[1] * A09[8] + ) + * h, + ( + K00[2] * A09[0] + + K01[2] * A09[1] + + K02[2] * A09[2] + + K03[2] * A09[3] + + K04[2] * A09[4] + + K05[2] * A09[5] + + K06[2] * A09[6] + + K07[2] * A09[7] + + K08[2] * A09[8] + ) + * h, + ) + dv = ( + ( + K00[3] * A09[0] + + K01[3] * A09[1] + + K02[3] * A09[2] + + K03[3] * A09[3] + + K04[3] * A09[4] + + K05[3] * A09[5] + + K06[3] * A09[6] + + K07[3] * A09[7] + + K08[3] * A09[8] + ) + * h, + ( + K00[4] * A09[0] + + K01[4] * A09[1] + + K02[4] * A09[2] + + K03[4] * A09[3] + + K04[4] * A09[4] + + K05[4] * A09[5] + + K06[4] * A09[6] + + K07[4] * A09[7] + + K08[4] * A09[8] + ) + * h, + ( + K00[5] * A09[0] + + K01[5] * A09[1] + + K02[5] * A09[2] + + K03[5] * A09[3] + + K04[5] * A09[4] + + K05[5] * A09[5] + + K06[5] * A09[6] + + K07[5] * A09[7] + + K08[5] * A09[8] + ) + * h, + ) + fr, fv = fun( + t + C[9] * h, + add_VV_hf(rr, dr), + add_VV_hf(vv, dv), + argk, + ) + K09 = *fr, *fv - dy = np.array( - [ - ( - K00[0] * B[0] - + K01[0] * B[1] - + K02[0] * B[2] - + K03[0] * B[3] - + K04[0] * B[4] - + K05[0] * B[5] - + K06[0] * B[6] - + K07[0] * B[7] - + K08[0] * B[8] - + K09[0] * B[9] - + K10[0] * B[10] - + K11[0] * B[11] - ) - * h, - ( - K00[1] * B[0] - + K01[1] * B[1] - + K02[1] * B[2] - + K03[1] * B[3] - + K04[1] * B[4] - + K05[1] * B[5] - + K06[1] * B[6] - + K07[1] * B[7] - + K08[1] * B[8] - + K09[1] * B[9] - + K10[1] * B[10] - + K11[1] * B[11] - ) - * h, - ( - K00[2] * B[0] - + K01[2] * B[1] - + K02[2] * B[2] - + K03[2] * B[3] - + K04[2] * B[4] - + K05[2] * B[5] - + K06[2] * B[6] - + K07[2] * B[7] - + K08[2] * B[8] - + K09[2] * B[9] - + K10[2] * B[10] - + K11[2] * B[11] - ) - * h, - ( - K00[3] * B[0] - + K01[3] * B[1] - + K02[3] * B[2] - + K03[3] * B[3] - + K04[3] * B[4] - + K05[3] * B[5] - + K06[3] * B[6] - + K07[3] * B[7] - + K08[3] * B[8] - + K09[3] * B[9] - + K10[3] * B[10] - + K11[3] * B[11] - ) - * h, - ( - K00[4] * B[0] - + K01[4] * B[1] - + K02[4] * B[2] - + K03[4] * B[3] - + K04[4] * B[4] - + K05[4] * B[5] - + K06[4] * B[6] - + K07[4] * B[7] - + K08[4] * B[8] - + K09[4] * B[9] - + K10[4] * B[10] - + K11[4] * B[11] - ) - * h, - ( - K00[5] * B[0] - + K01[4] * B[1] - + K02[5] * B[2] - + K03[5] * B[3] - + K04[5] * B[4] - + K05[5] * B[5] - + K06[5] * B[6] - + K07[5] * B[7] - + K08[5] * B[8] - + K09[5] * B[9] - + K10[5] * B[10] - + K11[5] * B[11] - ) - * h, - ] - ) - y_new = y + dy - f_new = fun(t + h, y_new, argk) + dr = ( + ( + K00[0] * A10[0] + + K01[0] * A10[1] + + K02[0] * A10[2] + + K03[0] * A10[3] + + K04[0] * A10[4] + + K05[0] * A10[5] + + K06[0] * A10[6] + + K07[0] * A10[7] + + K08[0] * A10[8] + + K09[0] * A10[9] + ) + * h, + ( + K00[1] * A10[0] + + K01[1] * A10[1] + + K02[1] * A10[2] + + K03[1] * A10[3] + + K04[1] * A10[4] + + K05[1] * A10[5] + + K06[1] * A10[6] + + K07[1] * A10[7] + + K08[1] * A10[8] + + K09[1] * A10[9] + ) + * h, + ( + K00[2] * A10[0] + + K01[2] * A10[1] + + K02[2] * A10[2] + + K03[2] * A10[3] + + K04[2] * A10[4] + + K05[2] * A10[5] + + K06[2] * A10[6] + + K07[2] * A10[7] + + K08[2] * A10[8] + + K09[2] * A10[9] + ) + * h, + ) + dv = ( + ( + K00[3] * A10[0] + + K01[3] * A10[1] + + K02[3] * A10[2] + + K03[3] * A10[3] + + K04[3] * A10[4] + + K05[3] * A10[5] + + K06[3] * A10[6] + + K07[3] * A10[7] + + K08[3] * A10[8] + + K09[3] * A10[9] + ) + * h, + ( + K00[4] * A10[0] + + K01[4] * A10[1] + + K02[4] * A10[2] + + K03[4] * A10[3] + + K04[4] * A10[4] + + K05[4] * A10[5] + + K06[4] * A10[6] + + K07[4] * A10[7] + + K08[4] * A10[8] + + K09[4] * A10[9] + ) + * h, + ( + K00[5] * A10[0] + + K01[5] * A10[1] + + K02[5] * A10[2] + + K03[5] * A10[3] + + K04[5] * A10[4] + + K05[5] * A10[5] + + K06[5] * A10[6] + + K07[5] * A10[7] + + K08[5] * A10[8] + + K09[5] * A10[9] + ) + * h, + ) + fr, fv = fun( + t + C[10] * h, + add_VV_hf(rr, dr), + add_VV_hf(vv, dv), + argk, + ) + K10 = *fr, *fv - K12 = f_new + dr = ( + ( + K00[0] * A11[0] + + K01[0] * A11[1] + + K02[0] * A11[2] + + K03[0] * A11[3] + + K04[0] * A11[4] + + K05[0] * A11[5] + + K06[0] * A11[6] + + K07[0] * A11[7] + + K08[0] * A11[8] + + K09[0] * A11[9] + + K10[0] * A11[10] + ) + * h, + ( + K00[1] * A11[0] + + K01[1] * A11[1] + + K02[1] * A11[2] + + K03[1] * A11[3] + + K04[1] * A11[4] + + K05[1] * A11[5] + + K06[1] * A11[6] + + K07[1] * A11[7] + + K08[1] * A11[8] + + K09[1] * A11[9] + + K10[1] * A11[10] + ) + * h, + ( + K00[2] * A11[0] + + K01[2] * A11[1] + + K02[2] * A11[2] + + K03[2] * A11[3] + + K04[2] * A11[4] + + K05[2] * A11[5] + + K06[2] * A11[6] + + K07[2] * A11[7] + + K08[2] * A11[8] + + K09[2] * A11[9] + + K10[2] * A11[10] + ) + * h, + ) + dv = ( + ( + K00[3] * A11[0] + + K01[3] * A11[1] + + K02[3] * A11[2] + + K03[3] * A11[3] + + K04[3] * A11[4] + + K05[3] * A11[5] + + K06[3] * A11[6] + + K07[3] * A11[7] + + K08[3] * A11[8] + + K09[3] * A11[9] + + K10[3] * A11[10] + ) + * h, + ( + K00[4] * A11[0] + + K01[4] * A11[1] + + K02[4] * A11[2] + + K03[4] * A11[3] + + K04[4] * A11[4] + + K05[4] * A11[5] + + K06[4] * A11[6] + + K07[4] * A11[7] + + K08[4] * A11[8] + + K09[4] * A11[9] + + K10[4] * A11[10] + ) + * h, + ( + K00[5] * A11[0] + + K01[5] * A11[1] + + K02[5] * A11[2] + + K03[5] * A11[3] + + K04[5] * A11[4] + + K05[5] * A11[5] + + K06[5] * A11[6] + + K07[5] * A11[7] + + K08[5] * A11[8] + + K09[5] * A11[9] + + K10[5] * A11[10] + ) + * h, + ) + fr, fv = fun( + t + C[11] * h, + add_VV_hf(rr, dr), + add_VV_hf(vv, dv), + argk, + ) + K11 = *fr, *fv - assert y_new.shape == (N_RV,) - assert f_new.shape == (N_RV,) + dr = ( + ( + K00[0] * B[0] + + K01[0] * B[1] + + K02[0] * B[2] + + K03[0] * B[3] + + K04[0] * B[4] + + K05[0] * B[5] + + K06[0] * B[6] + + K07[0] * B[7] + + K08[0] * B[8] + + K09[0] * B[9] + + K10[0] * B[10] + + K11[0] * B[11] + ) + * h, + ( + K00[1] * B[0] + + K01[1] * B[1] + + K02[1] * B[2] + + K03[1] * B[3] + + K04[1] * B[4] + + K05[1] * B[5] + + K06[1] * B[6] + + K07[1] * B[7] + + K08[1] * B[8] + + K09[1] * B[9] + + K10[1] * B[10] + + K11[1] * B[11] + ) + * h, + ( + K00[2] * B[0] + + K01[2] * B[1] + + K02[2] * B[2] + + K03[2] * B[3] + + K04[2] * B[4] + + K05[2] * B[5] + + K06[2] * B[6] + + K07[2] * B[7] + + K08[2] * B[8] + + K09[2] * B[9] + + K10[2] * B[10] + + K11[2] * B[11] + ) + * h, + ) + dv = ( + ( + K00[3] * B[0] + + K01[3] * B[1] + + K02[3] * B[2] + + K03[3] * B[3] + + K04[3] * B[4] + + K05[3] * B[5] + + K06[3] * B[6] + + K07[3] * B[7] + + K08[3] * B[8] + + K09[3] * B[9] + + K10[3] * B[10] + + K11[3] * B[11] + ) + * h, + ( + K00[4] * B[0] + + K01[4] * B[1] + + K02[4] * B[2] + + K03[4] * B[3] + + K04[4] * B[4] + + K05[4] * B[5] + + K06[4] * B[6] + + K07[4] * B[7] + + K08[4] * B[8] + + K09[4] * B[9] + + K10[4] * B[10] + + K11[4] * B[11] + ) + * h, + ( + K00[5] * B[0] + + K01[4] * B[1] + + K02[5] * B[2] + + K03[5] * B[3] + + K04[5] * B[4] + + K05[5] * B[5] + + K06[5] * B[6] + + K07[5] * B[7] + + K08[5] * B[8] + + K09[5] * B[9] + + K10[5] * B[10] + + K11[5] * B[11] + ) + * h, + ) + rr_new = add_VV_hf(rr, dr) + vv_new = add_VV_hf(vv, dr) + fr_new, fv_new = fun(t + h, rr_new, vv_new, argk) + K12 = *fr_new, *fv_new return ( - y_new, - f_new, + rr_new, + vv_new, + fr_new, + fv_new, ( K00, K01, diff --git a/src/hapsira/core/propagation/__init__.py b/src/hapsira/core/propagation/__init__.py index fb4bc643d..c7bda8291 100644 --- a/src/hapsira/core/propagation/__init__.py +++ b/src/hapsira/core/propagation/__init__.py @@ -1,36 +1 @@ """Low level propagation algorithms.""" - -from hapsira.core.propagation.base import func_twobody -from hapsira.core.propagation.cowell import cowell - -# from hapsira.core.propagation.danby import danby, danby_coe -# from hapsira.core.propagation.farnocchia import ( -# farnocchia_coe, -# farnocchia_rv as farnocchia, -# ) -# from hapsira.core.propagation.gooding import gooding, gooding_coe -# from hapsira.core.propagation.markley import markley, markley_coe -# from hapsira.core.propagation.mikkola import mikkola, mikkola_coe -# from hapsira.core.propagation.pimienta import pimienta, pimienta_coe -# from hapsira.core.propagation.recseries import recseries, recseries_coe -# from hapsira.core.propagation.vallado import vallado - -__all__ = [ - "cowell", - "func_twobody", - # "farnocchia_coe", - # "farnocchia", - # "vallado", - # "mikkola_coe", - # "mikkola", - # "markley_coe", - # "markley", - # "pimienta_coe", - # "pimienta", - # "gooding_coe", - # "gooding", - # "danby_coe", - # "danby", - # "recseries_coe", - # "recseries", -] diff --git a/src/hapsira/core/propagation/base.py b/src/hapsira/core/propagation/base.py index 5d289a5cb..fb3ee20b8 100644 --- a/src/hapsira/core/propagation/base.py +++ b/src/hapsira/core/propagation/base.py @@ -1,23 +1,29 @@ -from numba import njit as jit -import numpy as np +from ..jit import djit -@jit -def func_twobody(t0, u_, k): +__all__ = [ + "func_twobody_hf", +] + + +@djit +def func_twobody_hf(t0, rr, vv, k): """Differential equation for the initial value two body problem. Parameters ---------- t0 : float Time. - u_ : numpy.ndarray - Six component state vector [x, y, z, vx, vy, vz] (km, km/s). + rr : tuple[float,float,float] + Position vector + vv : tuple[float,float,float] + Velocity vector. k : float Standard gravitational parameter. """ - x, y, z, vx, vy, vz = u_ + x, y, z = rr + vx, vy, vz = vv r3 = (x**2 + y**2 + z**2) ** 1.5 - du = np.array([vx, vy, vz, -k * x / r3, -k * y / r3, -k * z / r3]) - return du + return (vx, vy, vz), (-k * x / r3, -k * y / r3, -k * z / r3) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 3615d586c..6a03adb46 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -1,20 +1,15 @@ import numpy as np -from ..jit import hjit from ..math.ivp import solve_ivp -from ..propagation.base import func_twobody +from ..propagation.base import func_twobody_hf -def cowelljit(func): - """ - Wrapper for hjit to track funcs for cowell - """ - compiled = hjit("Tuple([V,V])(f,V,V,f)")(func) - compiled.cowell = None # for debugging - return compiled +__all__ = [ + "cowell", +] -def cowell(k, r, v, tofs, rtol=1e-11, events=None, f=func_twobody): +def cowell(k, r, v, tofs, rtol=1e-11, events=None, f=func_twobody_hf): """ Scalar cowell @@ -24,7 +19,7 @@ def cowell(k, r, v, tofs, rtol=1e-11, events=None, f=func_twobody): tofs : ??? rtol : float ... or also ndarray? """ - # assert hasattr(f, "cowell") + assert hasattr(f, "djit") # DEBUG check for compiler flag assert isinstance(rtol, float) x, y, z = r diff --git a/src/hapsira/earth/__init__.py b/src/hapsira/earth/__init__.py index 6fb8e8144..3cd1af6c0 100644 --- a/src/hapsira/earth/__init__.py +++ b/src/hapsira/earth/__init__.py @@ -1,14 +1,13 @@ """Earth focused orbital mechanics routines.""" -from typing import Dict from astropy import units as u -import numpy as np from hapsira.bodies import Earth -from hapsira.core.jit import array_to_V_hf +from hapsira.core.jit import hjit, djit +from hapsira.core.math.linalg import add_VV_hf from hapsira.core.perturbations import J2_perturbation_hf -from hapsira.core.propagation import func_twobody +from hapsira.core.propagation.base import func_twobody_hf from hapsira.earth.enums import EarthGravity from hapsira.twobody.propagation import CowellPropagator @@ -74,40 +73,34 @@ def propagate(self, tof, atmosphere=None, gravity=None, *args): A new EarthSatellite with the propagated Orbit """ - ad_kwargs: Dict[object, dict] = {} - perturbations: Dict[object, dict] = {} - - def ad(t0, state, k, perturbations): # TODO compile - rr, vv = array_to_V_hf(state[:3]), array_to_V_hf(state[3:]) - if perturbations: - return np.sum( - [ - f(t0=t0, rr=rr, vv=vv, k=k, **p) - for f, p in perturbations.items() - ], - axis=0, - ) - else: - return np.array([0, 0, 0]) - - if gravity is EarthGravity.J2: # TODO move into compiled `ad` function - perturbations[J2_perturbation_hf] = { - "J2": Earth.J2.value, - "R": Earth.R.to_value(u.km), - } + + if gravity not in (None, EarthGravity.J2): + raise NotImplementedError + if atmosphere is not None: # Cannot compute density without knowing the state, # the perturbations parameters are not always fixed - # TODO: This whole function probably needs a refactoring raise NotImplementedError - def f(t0, state, k): # TODO compile - du_kep = func_twobody(t0, state, k) - ax, ay, az = ad(t0, state, k, perturbations) - du_ad = np.array([0, 0, 0, ax, ay, az]) + if gravity: + J2_ = Earth.J2.value + R_ = Earth.R.to_value(u.km) + + @hjit("V(f,V,V,f)") + def ad_hf(t0, rr, vv, k): + return J2_perturbation_hf(t0, rr, vv, k, J2_, R_) + + else: + + @hjit("V(f,V,V,f)") + def ad_hf(t0, rr, vv, k): + return 0.0, 0.0, 0.0 - return du_kep + du_ad + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad_vv = ad_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad_vv) - ad_kwargs.update(perturbations=perturbations) - new_orbit = self.orbit.propagate(tof, method=CowellPropagator(f=f)) + new_orbit = self.orbit.propagate(tof, method=CowellPropagator(f=f_hf)) return EarthSatellite(new_orbit, self.spacecraft) diff --git a/src/hapsira/twobody/propagation/cowell.py b/src/hapsira/twobody/propagation/cowell.py index 693369ecb..dd98d7855 100644 --- a/src/hapsira/twobody/propagation/cowell.py +++ b/src/hapsira/twobody/propagation/cowell.py @@ -2,8 +2,8 @@ from astropy import units as u -from hapsira.core.propagation import cowell -from hapsira.core.propagation.base import func_twobody +from hapsira.core.propagation.cowell import cowell +from hapsira.core.propagation.base import func_twobody_hf from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import RVState @@ -27,7 +27,7 @@ class CowellPropagator: PropagatorKind.ELLIPTIC | PropagatorKind.PARABOLIC | PropagatorKind.HYPERBOLIC ) - def __init__(self, rtol=1e-11, events=None, f=func_twobody): + def __init__(self, rtol=1e-11, events=None, f=func_twobody_hf): self._rtol = rtol self._events = events self._f = f diff --git a/tests/tests_twobody/test_events.py b/tests/tests_twobody/test_events.py index b9fbaf093..b43baac3b 100644 --- a/tests/tests_twobody/test_events.py +++ b/tests/tests_twobody/test_events.py @@ -8,9 +8,10 @@ from hapsira.bodies import Earth from hapsira.constants import H0_earth, rho0_earth from hapsira.core.events import line_of_sight_gf -from hapsira.core.jit import array_to_V_hf +from hapsira.core.jit import djit +from hapsira.core.math.linalg import add_VV_hf from hapsira.core.perturbations import atmospheric_drag_exponential_hf -from hapsira.core.propagation import func_twobody +from hapsira.core.propagation.base import func_twobody_hf from hapsira.twobody import Orbit from hapsira.twobody.events import ( AltitudeCrossEvent, @@ -48,23 +49,23 @@ def test_altitude_crossing(): altitude_cross_event = AltitudeCrossEvent(thresh_alt, R) events = [altitude_cross_event] - def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = atmospheric_drag_exponential_hf( + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = atmospheric_drag_exponential_hf( t0, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), + rr, + vv, k, - R=R, - C_D=C_D, - A_over_m=A_over_m, - H0=H0, - rho0=rho0, + R, + C_D, + A_over_m, + H0, + rho0, ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - method = CowellPropagator(events=events, f=f) + method = CowellPropagator(events=events, f=f_hf) rr, _ = method.propagate_many( orbit._state, tofs, @@ -339,8 +340,8 @@ def test_line_of_sight(): r_sun = np.array([122233179, -76150708, 33016374]) << u.km R = Earth.R.to(u.km).value - los = line_of_sight_gf(r1.value, r2.value, R) - los_with_sun = line_of_sight_gf(r1.value, r_sun.value, R) + los = line_of_sight_gf(r1.value, r2.value, R) # pylint: disable=E1120 + los_with_sun = line_of_sight_gf(r1.value, r_sun.value, R) # pylint: disable=E1120 assert los < 0 # No LOS condition. assert los_with_sun >= 0 # LOS condition. diff --git a/tests/tests_twobody/test_perturbations.py b/tests/tests_twobody/test_perturbations.py index 09c6bdd94..e89b048a8 100644 --- a/tests/tests_twobody/test_perturbations.py +++ b/tests/tests_twobody/test_perturbations.py @@ -9,18 +9,19 @@ from hapsira.bodies import Earth, Moon, Sun from hapsira.constants import H0_earth, Wdivc_sun, rho0_earth from hapsira.core.elements import rv2coe_gf, RV2COE_TOL -from hapsira.core.jit import array_to_V_hf, hjit -from hapsira.core.math.linalg import mul_Vs_hf, norm_hf +from hapsira.core.jit import hjit, djit +from hapsira.core.math.linalg import add_VV_hf, mul_Vs_hf, norm_hf from hapsira.core.perturbations import ( # pylint: disable=E1120,E1136 J2_perturbation_hf, J3_perturbation_hf, - atmospheric_drag_hf, + # atmospheric_drag_hf, # TODO reactivate test atmospheric_drag_exponential_hf, radiation_pressure_hf, third_body_hf, ) -from hapsira.core.propagation import func_twobody -from hapsira.earth.atmosphere import COESA76 +from hapsira.core.propagation.base import func_twobody_hf + +# from hapsira.earth.atmosphere import COESA76 # TODO reactivate test from hapsira.ephem import build_ephem_interpolant from hapsira.twobody import Orbit from hapsira.twobody.events import LithobrakeEvent @@ -37,21 +38,23 @@ def test_J2_propagation_Earth(): orbit = Orbit.from_vectors(Earth, r0 * u.km, v0 * u.km / u.s) tofs = [48.0] * u.h + J2 = Earth.J2.value + R_ = Earth.R.to(u.km).value - def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = J2_perturbation_hf( + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = J2_perturbation_hf( t0, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), + rr, + vv, k, - J2=Earth.J2.value, - R=Earth.R.to(u.km).value, + J2, + R_, ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - method = CowellPropagator(f=f) + method = CowellPropagator(f=f_hf) rr, vv = method.propagate_many(orbit._state, tofs) k = Earth.k.to(u.km**3 / u.s**2).value @@ -126,51 +129,53 @@ def test_J3_propagation_Earth(test_params): nu=nu_ini, ) - def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = J2_perturbation_hf( + J2 = Earth.J2.value + R_ = Earth.R.to(u.km).value + + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = J2_perturbation_hf( t0, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), + rr, + vv, k, - J2=Earth.J2.value, - R=Earth.R.to(u.km).value, + J2, + R_, ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) tofs = np.linspace(0, 10.0 * u.day, 1000) - method = CowellPropagator(rtol=1e-8, f=f) + method = CowellPropagator(rtol=1e-8, f=f_hf) r_J2, v_J2 = method.propagate_many( orbit._state, tofs, ) - def f_combined(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = np.array( - J2_perturbation_hf( - t0, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), - k, - J2=Earth.J2.value, - R=Earth.R.to_value(u.km), - ) - ) + np.array( - J3_perturbation_hf( - t0, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), - k, - J3=Earth.J3.value, - R=Earth.R.to_value(u.km), - ) + J3 = Earth.J3.value + + @djit + def f_combined_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad_J2 = J2_perturbation_hf( + t0, + rr, + vv, + k, + J2, + R_, + ) + du_ad_J3 = J3_perturbation_hf( + t0, + rr, + vv, + k, + J3, + R_, ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_vv, add_VV_hf(du_ad_J2, du_ad_J3)) - method = CowellPropagator(rtol=1e-8, f=f_combined) + method = CowellPropagator(rtol=1e-8, f=f_combined_hf) r_J3, v_J3 = method.propagate_many( orbit._state, tofs, @@ -266,23 +271,23 @@ def test_atmospheric_drag_exponential(): # dr_expected = F_r * tof (Newton's integration formula), where # F_r = -B rho(r) |r|^2 sqrt(k / |r|^3) = -B rho(r) sqrt(k |r|) - def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = atmospheric_drag_exponential_hf( + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = atmospheric_drag_exponential_hf( t0, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), + rr, + vv, k, - R=R, - C_D=C_D, - A_over_m=A_over_m, - H0=H0, - rho0=rho0, + R, + C_D, + A_over_m, + H0, + rho0, ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - method = CowellPropagator(f=f) + method = CowellPropagator(f=f_hf) rr, _ = method.propagate_many( orbit._state, [tof] * u.s, @@ -316,23 +321,23 @@ def test_atmospheric_demise(): lithobrake_event = LithobrakeEvent(R) events = [lithobrake_event] - def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = atmospheric_drag_exponential_hf( + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = atmospheric_drag_exponential_hf( t0, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), + rr, + vv, k, - R=R, - C_D=C_D, - A_over_m=A_over_m, - H0=H0, - rho0=rho0, + R, + C_D, + A_over_m, + H0, + rho0, ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - method = CowellPropagator(events=events, f=f) + method = CowellPropagator(events=events, f=f_hf) rr, _ = method.propagate_many( orbit._state, tofs, @@ -347,7 +352,7 @@ def f(t0, u_, k): lithobrake_event = LithobrakeEvent(R) events = [lithobrake_event] - method = CowellPropagator(events=events, f=f) + method = CowellPropagator(events=events, f=f_hf) rr, _ = method.propagate_many( orbit._state, tofs, @@ -356,55 +361,60 @@ def f(t0, u_, k): assert lithobrake_event.last_t == tofs[-1] -@pytest.mark.slow -def test_atmospheric_demise_coesa76(): - # Test an orbital decay that hits Earth. No analytic solution. - R = Earth.R.to(u.km).value +# @pytest.mark.slow +# def test_atmospheric_demise_coesa76(): +# # Test an orbital decay that hits Earth. No analytic solution. +# R = Earth.R.to(u.km).value - orbit = Orbit.circular(Earth, 250 * u.km) - t_decay = 7.17 * u.d +# orbit = Orbit.circular(Earth, 250 * u.km) +# t_decay = 7.17 * u.d - # Parameters of a body - C_D = 2.2 # Dimensionless (any value would do) - A_over_m = ((np.pi / 4.0) * (u.m**2) / (100 * u.kg)).to_value( - u.km**2 / u.kg - ) # km^2/kg +# # Parameters of a body +# C_D = 2.2 # Dimensionless (any value would do) +# A_over_m = ((np.pi / 4.0) * (u.m**2) / (100 * u.kg)).to_value( +# u.km**2 / u.kg +# ) # km^2/kg - tofs = [365] * u.d +# tofs = [365] * u.d - lithobrake_event = LithobrakeEvent(R) - events = [lithobrake_event] +# lithobrake_event = LithobrakeEvent(R) +# events = [lithobrake_event] - coesa76 = COESA76() +# coesa76 = COESA76() - def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) +# @djit +# def f_hf(t0, rr, vv, k): +# du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) - # Avoid undershooting H below attractor radius R - H = max(norm(u_[:3]), R) - rho = coesa76.density((H - R) * u.km).to_value(u.kg / u.km**3) +# # Avoid undershooting H below attractor radius R +# H = norm_hf(rr) +# if H < R: +# H = R - ax, ay, az = atmospheric_drag_hf( - t0, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), - k, - C_D=C_D, - A_over_m=A_over_m, - rho=rho, - ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad +# rho = coesa76.density( +# (H - R) * u.km +# ).to_value(u.kg / u.km**3) # TODO jit'ed ... move to core?!? - method = CowellPropagator(events=events, f=f) - rr, _ = method.propagate_many( - orbit._state, - tofs, - ) +# du_ad = atmospheric_drag_hf( +# t0, +# rr, +# vv, +# k, +# C_D, +# A_over_m, +# rho, +# ) +# return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - assert_quantity_allclose(norm(rr[0].to(u.km).value), R, atol=1) # Below 1km +# method = CowellPropagator(events=events, f=f) +# rr, _ = method.propagate_many( +# orbit._state, +# tofs, +# ) - assert_quantity_allclose(lithobrake_event.last_t, t_decay, rtol=1e-2) +# assert_quantity_allclose(norm(rr[0].to(u.km).value), R, atol=1) # Below 1km + +# assert_quantity_allclose(lithobrake_event.last_t, t_decay, rtol=1e-2) @pytest.mark.slow @@ -430,18 +440,17 @@ def test_cowell_works_with_small_perturbations(): initial = Orbit.from_vectors(Earth, r0, v0) - def accel(t0, state, k): - v_vec = state[3:] - norm_v = (v_vec * v_vec).sum() ** 0.5 - return 1e-5 * v_vec / norm_v + @hjit("V(f,V,V,f)") + def accel_hf(t0, rr, vv, k): + return mul_Vs_hf(vv, 1e-5 / norm_hf(vv)) - def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = accel(t0, u_, k) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = accel_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - final = initial.propagate(3 * u.day, method=CowellPropagator(f=f)) + final = initial.propagate(3 * u.day, method=CowellPropagator(f=f_hf)) # TODO: Accuracy reduced after refactor, # but unclear what are we comparing against @@ -456,18 +465,18 @@ def test_cowell_converges_with_small_perturbations(): initial = Orbit.from_vectors(Earth, r0, v0) - def accel(t0, state, k): - v_vec = state[3:] - norm_v = (v_vec * v_vec).sum() ** 0.5 - return 0.0 * v_vec / norm_v + @hjit("V(f,V,V,f)") + def accel_hf(t0, rr, vv, k): + norm_v = norm_hf(vv) + return mul_Vs_hf(vv, 0.0 / norm_v) - def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = accel(t0, u_, k) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = accel_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - final = initial.propagate(initial.period, method=CowellPropagator(f=f)) + final = initial.propagate(initial.period, method=CowellPropagator(f=f_hf)) assert_quantity_allclose(final.r, initial.r) assert_quantity_allclose(final.v, initial.v) @@ -608,21 +617,22 @@ def test_3rd_body_Curtis(test_params): end=epoch + test_params["tof"], ) body_r = build_ephem_interpolant(body, body_epochs) + k_third = body.k.to_value(u.km**3 / u.s**2) - def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = third_body_hf( + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = third_body_hf( t0, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), + rr, + vv, k, - body.k.to_value(u.km**3 / u.s**2), # k_third + k_third, body_r, # perturbation_body ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - method = CowellPropagator(rtol=1e-10, f=f) + method = CowellPropagator(rtol=1e-10, f=f_hf) rr, vv = method.propagate_many( initial._state, np.linspace(0, tof, 400) << u.s, @@ -704,25 +714,28 @@ def sun_normalized_hf(t0): r = sun_r(t0) # sun_r is hf, returns V return mul_Vs_hf(r, 149600000 / norm_hf(r)) - def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = radiation_pressure_hf( + R_ = Earth.R.to(u.km).value + Wdivc_s = Wdivc_sun.value + + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = radiation_pressure_hf( t0, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), + rr, + vv, k, - R=Earth.R.to(u.km).value, - C_R=2.0, - A_over_m=2e-4 / 100, - Wdivc_s=Wdivc_sun.value, - star=sun_normalized_hf, + R_, + 2.0, # C_R + 2e-4 / 100, # A_over_m + Wdivc_s, + sun_normalized_hf, # star ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) method = CowellPropagator( rtol=1e-8, - f=f, + f=f_hf, ) rr, vv = method.propagate_many( initial._state, diff --git a/tests/tests_twobody/test_propagation.py b/tests/tests_twobody/test_propagation.py index 1a8a4bdad..cd3dfd5e1 100644 --- a/tests/tests_twobody/test_propagation.py +++ b/tests/tests_twobody/test_propagation.py @@ -10,7 +10,9 @@ from hapsira.bodies import Earth, Moon, Sun from hapsira.constants import J2000 from hapsira.core.elements import rv2coe_gf, RV2COE_TOL -from hapsira.core.propagation import func_twobody +from hapsira.core.jit import djit, hjit +from hapsira.core.math.linalg import add_VV_hf, mul_Vs_hf, norm_hf +from hapsira.core.propagation.base import func_twobody_hf from hapsira.examples import iss from hapsira.frames import Planes from hapsira.twobody import Orbit @@ -289,22 +291,21 @@ def test_cowell_propagation_circle_to_circle(): # From [Edelbaum, 1961] accel = 1e-7 - def constant_accel(t0, u_, k): - v = u_[3:] - norm_v = (v[0] ** 2 + v[1] ** 2 + v[2] ** 2) ** 0.5 - return accel * v / norm_v + @hjit("V(f,V,V,f)") + def constant_accel_hf(t0, rr, vv, k): + norm_v = norm_hf(vv) + return mul_Vs_hf(vv, accel / norm_v) - def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = constant_accel(t0, u_, k) - du_ad = np.array([0, 0, 0, ax, ay, az]) - - return du_kep + du_ad + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = constant_accel_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) ss = Orbit.circular(Earth, 500 * u.km) tofs = [20] * ss.period - method = CowellPropagator(f=f) + method = CowellPropagator(f=f_hf) rrs, vvs = method.propagate_many(ss._state, tofs) orb_final = Orbit.from_vectors(Earth, rrs[0], vvs[0]) diff --git a/tests/tests_twobody/test_thrust.py b/tests/tests_twobody/test_thrust.py index bf59399e4..93f2300f4 100644 --- a/tests/tests_twobody/test_thrust.py +++ b/tests/tests_twobody/test_thrust.py @@ -4,8 +4,9 @@ import pytest from hapsira.bodies import Earth -from hapsira.core.jit import array_to_V_hf -from hapsira.core.propagation import func_twobody +from hapsira.core.jit import djit +from hapsira.core.math.linalg import add_VV_hf +from hapsira.core.propagation.base import func_twobody_hf from hapsira.core.thrust.change_a_inc import change_a_inc_hb from hapsira.core.thrust.change_argp import change_argp_hb from hapsira.core.thrust.change_ecc_inc import beta_vf as beta_change_ecc_inc @@ -33,19 +34,19 @@ def test_leo_geo_numerical_safe(inc_0): k = Earth.k.to(u.km**3 / u.s**2) - a_d, _, t_f = change_a_inc(k, a_0, a_f, inc_0, inc_f, f) + a_d_hf, _, t_f = change_a_inc(k, a_0, a_f, inc_0, inc_f, f) # Retrieve r and v from initial orbit s0 = Orbit.circular(Earth, a_0 - Earth.R, inc_0) # Propagate orbit - def f_leo_geo(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + @djit + def f_leo_geo_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = a_d_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - sf = s0.propagate(t_f, method=CowellPropagator(rtol=1e-6, f=f_leo_geo)) + sf = s0.propagate(t_f, method=CowellPropagator(rtol=1e-6, f=f_leo_geo_hf)) assert_allclose(sf.a.to(u.km).value, a_f.value, rtol=1e-3) assert_allclose(sf.ecc.value, 0.0, atol=1e-2) @@ -71,13 +72,13 @@ def test_leo_geo_numerical_fast(inc_0): s0 = Orbit.circular(Earth, a_0 * u.km - Earth.R, inc_0 * u.rad) # Propagate orbit - def f_leo_geo(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d_hf(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + @djit + def f_leo_geo_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = a_d_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - sf = s0.propagate(t_f * u.s, method=CowellPropagator(rtol=1e-6, f=f_leo_geo)) + sf = s0.propagate(t_f * u.s, method=CowellPropagator(rtol=1e-6, f=f_leo_geo_hf)) assert_allclose(sf.a.to(u.km).value, a_f, rtol=1e-3) assert_allclose(sf.ecc.value, 0.0, atol=1e-2) @@ -130,13 +131,15 @@ def test_sso_disposal_numerical(ecc_0, ecc_f): a_d_hf, _, t_f = change_ecc_quasioptimal(s0, ecc_f, f) # Propagate orbit - def f_ss0_disposal(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d_hf(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + @djit + def f_ss0_disposal_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = a_d_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - sf = s0.propagate(t_f * u.s, method=CowellPropagator(rtol=1e-8, f=f_ss0_disposal)) + sf = s0.propagate( + t_f * u.s, method=CowellPropagator(rtol=1e-8, f=f_ss0_disposal_hf) + ) assert_allclose(sf.ecc.value, ecc_f, rtol=1e-4, atol=1e-4) @@ -199,13 +202,13 @@ def test_geo_cases_numerical(ecc_0, ecc_f): a_d_hf, _, t_f = change_ecc_inc(orb_0=s0, ecc_f=ecc_f, inc_f=inc_f, f=f) # Propagate orbit - def f_geo(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d_hf(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + @djit + def f_geo_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = a_d_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - sf = s0.propagate(t_f, method=CowellPropagator(rtol=1e-8, f=f_geo)) + sf = s0.propagate(t_f, method=CowellPropagator(rtol=1e-8, f=f_geo_hf)) assert_allclose(sf.ecc.value, ecc_f, rtol=1e-2, atol=1e-2) assert_allclose(sf.inc.to_value(u.rad), inc_f.to_value(u.rad), rtol=1e-1) @@ -282,13 +285,13 @@ def test_soyuz_standard_gto_numerical_safe(): ) # Propagate orbit - def f_soyuz(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d_hf(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + @djit + def f_soyuz_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = a_d_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) - sf = s0.propagate(t_f, method=CowellPropagator(rtol=1e-8, f=f_soyuz)) + sf = s0.propagate(t_f, method=CowellPropagator(rtol=1e-8, f=f_soyuz_hf)) assert_allclose(sf.argp.to_value(u.rad), argp_f.to_value(u.rad), rtol=1e-4) @@ -320,15 +323,15 @@ def test_soyuz_standard_gto_numerical_fast(): ) # Propagate orbit - def f_soyuz(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = a_d_hf(t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + @djit + def f_soyuz_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + du_ad = a_d_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) sf = s0.propagate( t_f * u.s, - method=CowellPropagator(rtol=1e-8, f=f_soyuz), + method=CowellPropagator(rtol=1e-8, f=f_soyuz_hf), ) assert_allclose(sf.argp.to(u.rad).value, argp_f, rtol=1e-4) From ab905360695e6161aac7699e1aa809d8f80b42cf Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 22 Jan 2024 20:35:06 +0100 Subject: [PATCH 151/346] compiles, tests break --- src/hapsira/core/math/ivp/_rkstep.py | 31 ++++++++++++++-------------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index 3ac459caa..2f1bc0cb4 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -1,7 +1,9 @@ from typing import Callable, Tuple +from numpy import float32 as f4 + from ._dop853_coefficients import A as _A, B as _B, C as _C -from ...math.linalg import add_VV_hf +from ..linalg import add_VV_hf from ...jit import hjit, DSIG __all__ = [ @@ -13,20 +15,19 @@ N_RV = 6 N_STAGES = 12 -# A = tuple(tuple(float(number) for number in line) for line in A[:N_STAGES, :N_STAGES]) -A01 = tuple(float(number) for number in _A[1, :N_STAGES]) -A02 = tuple(float(number) for number in _A[2, :N_STAGES]) -A03 = tuple(float(number) for number in _A[3, :N_STAGES]) -A04 = tuple(float(number) for number in _A[4, :N_STAGES]) -A05 = tuple(float(number) for number in _A[5, :N_STAGES]) -A06 = tuple(float(number) for number in _A[6, :N_STAGES]) -A07 = tuple(float(number) for number in _A[7, :N_STAGES]) -A08 = tuple(float(number) for number in _A[8, :N_STAGES]) -A09 = tuple(float(number) for number in _A[9, :N_STAGES]) -A10 = tuple(float(number) for number in _A[10, :N_STAGES]) -A11 = tuple(float(number) for number in _A[11, :N_STAGES]) -B = tuple(float(number) for number in _B) -C = tuple(float(number) for number in _C[:N_STAGES]) +A01 = tuple(f4(number) for number in _A[1, :N_STAGES]) +A02 = tuple(f4(number) for number in _A[2, :N_STAGES]) +A03 = tuple(f4(number) for number in _A[3, :N_STAGES]) +A04 = tuple(f4(number) for number in _A[4, :N_STAGES]) +A05 = tuple(f4(number) for number in _A[5, :N_STAGES]) +A06 = tuple(f4(number) for number in _A[6, :N_STAGES]) +A07 = tuple(f4(number) for number in _A[7, :N_STAGES]) +A08 = tuple(f4(number) for number in _A[8, :N_STAGES]) +A09 = tuple(f4(number) for number in _A[9, :N_STAGES]) +A10 = tuple(f4(number) for number in _A[10, :N_STAGES]) +A11 = tuple(f4(number) for number in _A[11, :N_STAGES]) +B = tuple(f4(number) for number in _B) +C = tuple(f4(number) for number in _C[:N_STAGES]) _KSIG = ( "Tuple([" From f36157c73a23505d5dfc7db79149507643a87fce Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 22 Jan 2024 20:57:55 +0100 Subject: [PATCH 152/346] fix var name --- src/hapsira/core/math/ivp/_rkstep.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index 2f1bc0cb4..48762bc9b 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -816,7 +816,7 @@ def rk_step_hf( * h, ) rr_new = add_VV_hf(rr, dr) - vv_new = add_VV_hf(vv, dr) + vv_new = add_VV_hf(vv, dv) fr_new, fv_new = fun(t + h, rr_new, vv_new, argk) K12 = *fr_new, *fv_new From a7038ef6a478f2fe8e91e1ed1315e18550037043 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 23 Jan 2024 10:39:07 +0100 Subject: [PATCH 153/346] fix DOP853 by forcing f8; make PRECISION a setting --- src/hapsira/core/jit.py | 14 +++++++--- src/hapsira/core/math/ivp/_rkstep.py | 39 +++++++++++++++++----------- src/hapsira/settings.py | 7 +++++ 3 files changed, 41 insertions(+), 19 deletions(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 29a2da490..79d702e6b 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -9,7 +9,6 @@ __all__ = [ - "PRECISIONS", "DSIG", "hjit", "vjit", @@ -26,7 +25,12 @@ logger.debug("jit inline: %s", "yes" if settings["INLINE"].value else "no") logger.debug("jit nopython: %s", "yes" if settings["NOPYTHON"].value else "no") -PRECISIONS = ("f4", "f8") # TODO allow f2, i.e. half, for CUDA at least? +_PRECISIONS = ( + settings["PRECISION"].value, +) # TODO again allow to compile for multiple precision? +logger.debug("jit precision: %s", settings["PRECISION"].value) +if settings["PRECISION"].value != "f8": + logger.warning("jit precision: DOP853 as used by Cowell's method requires f8!") DSIG = "Tuple([V,V])(f,V,V,f)" @@ -56,7 +60,9 @@ def _parse_signatures(signature: str, noreturn: bool = False) -> Union[str, List ) return signature - if any(level in signature for level in PRECISIONS): # leave this signature as it is + if any( + level in signature for level in _PRECISIONS + ): # leave this signature as it is logger.warning( "jit signature: precision specified, not parsing (%s)", signature ) @@ -72,7 +78,7 @@ def _parse_signatures(signature: str, noreturn: bool = False) -> Union[str, List "F", "FunctionType" ) # TODO does not work for CUDA yet - return [signature.replace("f", dtype) for dtype in PRECISIONS] + return [signature.replace("f", dtype) for dtype in _PRECISIONS] def hjit(*args, **kwargs) -> Callable: diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index 48762bc9b..178745b5b 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -1,10 +1,19 @@ from typing import Callable, Tuple -from numpy import float32 as f4 - from ._dop853_coefficients import A as _A, B as _B, C as _C from ..linalg import add_VV_hf from ...jit import hjit, DSIG +from ....settings import settings + +if settings["PRECISION"].value == "f8": + from numpy import float64 as float_ +elif settings["PRECISION"].value == "f4": + from numpy import float32 as float_ +elif settings["PRECISION"].value == "f2": + from numpy import float16 as float_ +else: + raise ValueError("unsupported precision") + __all__ = [ "rk_step_hf", @@ -15,19 +24,19 @@ N_RV = 6 N_STAGES = 12 -A01 = tuple(f4(number) for number in _A[1, :N_STAGES]) -A02 = tuple(f4(number) for number in _A[2, :N_STAGES]) -A03 = tuple(f4(number) for number in _A[3, :N_STAGES]) -A04 = tuple(f4(number) for number in _A[4, :N_STAGES]) -A05 = tuple(f4(number) for number in _A[5, :N_STAGES]) -A06 = tuple(f4(number) for number in _A[6, :N_STAGES]) -A07 = tuple(f4(number) for number in _A[7, :N_STAGES]) -A08 = tuple(f4(number) for number in _A[8, :N_STAGES]) -A09 = tuple(f4(number) for number in _A[9, :N_STAGES]) -A10 = tuple(f4(number) for number in _A[10, :N_STAGES]) -A11 = tuple(f4(number) for number in _A[11, :N_STAGES]) -B = tuple(f4(number) for number in _B) -C = tuple(f4(number) for number in _C[:N_STAGES]) +A01 = tuple(float_(number) for number in _A[1, :N_STAGES]) +A02 = tuple(float_(number) for number in _A[2, :N_STAGES]) +A03 = tuple(float_(number) for number in _A[3, :N_STAGES]) +A04 = tuple(float_(number) for number in _A[4, :N_STAGES]) +A05 = tuple(float_(number) for number in _A[5, :N_STAGES]) +A06 = tuple(float_(number) for number in _A[6, :N_STAGES]) +A07 = tuple(float_(number) for number in _A[7, :N_STAGES]) +A08 = tuple(float_(number) for number in _A[8, :N_STAGES]) +A09 = tuple(float_(number) for number in _A[9, :N_STAGES]) +A10 = tuple(float_(number) for number in _A[10, :N_STAGES]) +A11 = tuple(float_(number) for number in _A[11, :N_STAGES]) +B = tuple(float_(number) for number in _B) +C = tuple(float_(number) for number in _C[:N_STAGES]) _KSIG = ( "Tuple([" diff --git a/src/hapsira/settings.py b/src/hapsira/settings.py index f12a4a881..a18e06060 100644 --- a/src/hapsira/settings.py +++ b/src/hapsira/settings.py @@ -133,6 +133,13 @@ def __init__(self): True, ) ) + self._add( + Setting( + "PRECISION", + "f8", + options=("f2", "f4", "f8"), + ) + ) def _add(self, setting: Setting): """ From 3393fdf179b25631217c937f0c93c3ff1abb64c9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 23 Jan 2024 16:48:36 +0100 Subject: [PATCH 154/346] consolidate coesa core code, jited --- .../{earth_atmosphere => earth}/__init__.py | 0 src/hapsira/core/earth/atmosphere/__init__.py | 0 .../util.py => earth/atmosphere/coesa.py} | 37 ++++++++++--------- .../atmosphere}/jacchia.py | 0 src/hapsira/earth/atmosphere/base.py | 10 ++--- src/hapsira/earth/atmosphere/jacchia.py | 2 +- 6 files changed, 26 insertions(+), 23 deletions(-) rename src/hapsira/core/{earth_atmosphere => earth}/__init__.py (100%) create mode 100644 src/hapsira/core/earth/atmosphere/__init__.py rename src/hapsira/core/{earth_atmosphere/util.py => earth/atmosphere/coesa.py} (75%) rename src/hapsira/core/{earth_atmosphere => earth/atmosphere}/jacchia.py (100%) diff --git a/src/hapsira/core/earth_atmosphere/__init__.py b/src/hapsira/core/earth/__init__.py similarity index 100% rename from src/hapsira/core/earth_atmosphere/__init__.py rename to src/hapsira/core/earth/__init__.py diff --git a/src/hapsira/core/earth/atmosphere/__init__.py b/src/hapsira/core/earth/atmosphere/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/hapsira/core/earth_atmosphere/util.py b/src/hapsira/core/earth/atmosphere/coesa.py similarity index 75% rename from src/hapsira/core/earth_atmosphere/util.py rename to src/hapsira/core/earth/atmosphere/coesa.py index b940e4e84..5c8d60315 100644 --- a/src/hapsira/core/earth_atmosphere/util.py +++ b/src/hapsira/core/earth/atmosphere/coesa.py @@ -1,10 +1,15 @@ """This script holds several utilities related to atmospheric computations.""" -from numba import njit as jit +from ...jit import hjit +__all__ = [ + "get_index_hf", + "check_altitude_hf", +] -@jit -def geometric_to_geopotential(z, r0): + +@hjit("f(f,f)") +def _geometric_to_geopotential_hf(z, r0): """Converts from given geometric altitude to geopotential one. Parameters @@ -23,11 +28,8 @@ def geometric_to_geopotential(z, r0): return h -z_to_h = geometric_to_geopotential - - -@jit -def geopotential_to_geometric(h, r0): +@hjit("f(f,f)") +def _geopotential_to_geometric_hf(h, r0): """Converts from given geopotential altitude to geometric one. Parameters @@ -46,11 +48,12 @@ def geopotential_to_geometric(h, r0): return z -h_to_z = geopotential_to_geometric +_z_to_h_hf = _geometric_to_geopotential_hf +_h_to_z_hf = _geopotential_to_geometric_hf -@jit -def gravity(z, g0, r0): +@hjit("f(f,f,f)") +def _gravity_hf(z, g0, r0): """Relates Earth gravity field magnitude with the geometric height. Parameters @@ -71,8 +74,8 @@ def gravity(z, g0, r0): return g -@jit -def _get_index(x, x_levels): +@hjit # ("i8(f,f)") # TODO use tuple with fixed length +def get_index_hf(x, x_levels): """Finds element in list and returns index. Parameters @@ -97,14 +100,14 @@ def _get_index(x, x_levels): return i - 1 -@jit -def _check_altitude(alt, r0, geometric): +@hjit("Tuple([f,f])(f,f,b1)") +def check_altitude_hf(alt, r0, geometric): # Get geometric and geopotential altitudes if geometric: z = alt - h = z_to_h(z, r0) + h = _z_to_h_hf(z, r0) else: h = alt - z = h_to_z(h, r0) + z = _h_to_z_hf(h, r0) return z, h diff --git a/src/hapsira/core/earth_atmosphere/jacchia.py b/src/hapsira/core/earth/atmosphere/jacchia.py similarity index 100% rename from src/hapsira/core/earth_atmosphere/jacchia.py rename to src/hapsira/core/earth/atmosphere/jacchia.py diff --git a/src/hapsira/earth/atmosphere/base.py b/src/hapsira/earth/atmosphere/base.py index dc5eeb2f7..2f81f9631 100644 --- a/src/hapsira/earth/atmosphere/base.py +++ b/src/hapsira/earth/atmosphere/base.py @@ -2,9 +2,9 @@ import astropy.units as u -from hapsira.core.earth_atmosphere.util import ( - _check_altitude as _check_altitude_fast, - _get_index as _get_index_fast, +from hapsira.core.earth.atmosphere.coesa import ( + check_altitude_hf, + get_index_hf, ) @@ -69,7 +69,7 @@ def _check_altitude(self, alt, r0, geometric=True): alt = alt.to_value(u.km) r0 = r0.to_value(u.km) - z, h = _check_altitude_fast(alt, r0, geometric) + z, h = check_altitude_hf(alt, r0, geometric) # TODO call from compiled context z, h = z * u.km, h * u.km # Assert in range @@ -98,5 +98,5 @@ def _get_index(self, x, x_levels): """ x = x.to_value(u.km) x_levels = (x_levels << u.km).value - i = _get_index_fast(x, x_levels) + i = get_index_hf(x, x_levels) # TODO call from compiled context return i diff --git a/src/hapsira/earth/atmosphere/jacchia.py b/src/hapsira/earth/atmosphere/jacchia.py index c2695248b..fcfb6e733 100644 --- a/src/hapsira/earth/atmosphere/jacchia.py +++ b/src/hapsira/earth/atmosphere/jacchia.py @@ -1,7 +1,7 @@ from astropy import units as u import numpy as np -from hapsira.core.earth_atmosphere.jacchia import ( +from hapsira.core.earth.atmosphere.jacchia import ( _altitude_profile as _altitude_profile_fast, _H_correction as _H_correction_fast, _O_and_O2_correction as _O_and_O2_correction_fast, From 4382f6f06baaaad34346040a68a699a806dca5b9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 23 Jan 2024 20:22:13 +0100 Subject: [PATCH 155/346] mv atmospheric data into core, keep units in high level api --- src/hapsira/core/earth/atmosphere/coesa76.py | 77 +++++++++++++++ .../earth/atmosphere/data/coesa62.dat | 0 .../earth/atmosphere/data/coesa76.dat | 0 .../earth/atmosphere/data/coesa76_p.dat | 0 .../earth/atmosphere/data/coesa76_rho.dat | 0 src/hapsira/earth/atmosphere/coesa62.py | 6 +- src/hapsira/earth/atmosphere/coesa76.py | 96 ++++++++++--------- .../tests_atmosphere/test_coesa76.py | 21 ++-- 8 files changed, 143 insertions(+), 57 deletions(-) create mode 100644 src/hapsira/core/earth/atmosphere/coesa76.py rename src/hapsira/{ => core}/earth/atmosphere/data/coesa62.dat (100%) rename src/hapsira/{ => core}/earth/atmosphere/data/coesa76.dat (100%) rename src/hapsira/{ => core}/earth/atmosphere/data/coesa76_p.dat (100%) rename src/hapsira/{ => core}/earth/atmosphere/data/coesa76_rho.dat (100%) diff --git a/src/hapsira/core/earth/atmosphere/coesa76.py b/src/hapsira/core/earth/atmosphere/coesa76.py new file mode 100644 index 000000000..7fc6e1b82 --- /dev/null +++ b/src/hapsira/core/earth/atmosphere/coesa76.py @@ -0,0 +1,77 @@ +from astropy.io import ascii as ascii_ +from astropy.utils.data import get_pkg_data_filename + +from numpy import float32 as f4 + +__all__ = [ + "R", + "R_air", + "k", + "Na", + "g0", + "r0", + "M0", + "P0", + "T0", + "Tinf", + "gamma", + "alpha", + "beta", + "S", + "b_levels", + "zb_levels", + "hb_levels", + "Tb_levels", + "Lb_levels", + "pb_levels", + "z_coeff", + "p_coeff", + "rho_coeff", +] + +# Following constants come from original U.S Atmosphere 1962 paper so a pure +# model of this atmosphere can be implemented +R = f4(8314.32) # u.J / u.kmol / u.K +R_air = f4(287.053) # u.J / u.kg / u.K +k = f4(1.380622e-23) # u.J / u.K +Na = f4(6.022169e-26) # 1 / u.kmol +g0 = f4(9.80665) # u.m / u.s**2 +r0 = f4(6356.766) # u.km +M0 = f4(28.9644) # u.kg / u.kmol +P0 = f4(101325) # u.Pa +T0 = f4(288.15) # u.K +Tinf = 1000 # u.K +gamma = f4(1.4) # one +alpha = f4(34.1632) # u.K / u.km +beta = f4(1.458e-6) # (u.kg / u.s / u.m / (u.K) ** 0.5) +S = f4(110.4) # u.K + +# Reading layer parameters file +coesa76_data = ascii_.read(get_pkg_data_filename("data/coesa76.dat")) +b_levels = tuple(f4(number) for number in coesa76_data["b"].data) +zb_levels = tuple(f4(number) for number in coesa76_data["Zb [km]"].data) # u.km +hb_levels = tuple(f4(number) for number in coesa76_data["Hb [km]"].data) # u.km +Tb_levels = tuple(f4(number) for number in coesa76_data["Tb [K]"].data) # u.K +Lb_levels = tuple(f4(number) for number in coesa76_data["Lb [K/km]"].data) # u.K / u.km +pb_levels = tuple(f4(number) for number in coesa76_data["pb [mbar]"].data) # u.mbar + +# Reading pressure and density coefficients files +p_data = ascii_.read(get_pkg_data_filename("data/coesa76_p.dat")) +rho_data = ascii_.read(get_pkg_data_filename("data/coesa76_rho.dat")) + +# Zip coefficients for each altitude +z_coeff = tuple(f4(number) for number in p_data["z [km]"].data) # u.km +p_coeff = ( + tuple(f4(number) for number in p_data["A"].data), + tuple(f4(number) for number in p_data["B"].data), + tuple(f4(number) for number in p_data["C"].data), + tuple(f4(number) for number in p_data["D"].data), + tuple(f4(number) for number in p_data["E"].data), +) +rho_coeff = ( + tuple(f4(number) for number in rho_data["A"].data), + tuple(f4(number) for number in rho_data["B"].data), + tuple(f4(number) for number in rho_data["C"].data), + tuple(f4(number) for number in rho_data["D"].data), + tuple(f4(number) for number in rho_data["E"].data), +) diff --git a/src/hapsira/earth/atmosphere/data/coesa62.dat b/src/hapsira/core/earth/atmosphere/data/coesa62.dat similarity index 100% rename from src/hapsira/earth/atmosphere/data/coesa62.dat rename to src/hapsira/core/earth/atmosphere/data/coesa62.dat diff --git a/src/hapsira/earth/atmosphere/data/coesa76.dat b/src/hapsira/core/earth/atmosphere/data/coesa76.dat similarity index 100% rename from src/hapsira/earth/atmosphere/data/coesa76.dat rename to src/hapsira/core/earth/atmosphere/data/coesa76.dat diff --git a/src/hapsira/earth/atmosphere/data/coesa76_p.dat b/src/hapsira/core/earth/atmosphere/data/coesa76_p.dat similarity index 100% rename from src/hapsira/earth/atmosphere/data/coesa76_p.dat rename to src/hapsira/core/earth/atmosphere/data/coesa76_p.dat diff --git a/src/hapsira/earth/atmosphere/data/coesa76_rho.dat b/src/hapsira/core/earth/atmosphere/data/coesa76_rho.dat similarity index 100% rename from src/hapsira/earth/atmosphere/data/coesa76_rho.dat rename to src/hapsira/core/earth/atmosphere/data/coesa76_rho.dat diff --git a/src/hapsira/earth/atmosphere/coesa62.py b/src/hapsira/earth/atmosphere/coesa62.py index 225772ea0..eee4ed245 100644 --- a/src/hapsira/earth/atmosphere/coesa62.py +++ b/src/hapsira/earth/atmosphere/coesa62.py @@ -53,7 +53,7 @@ """ from astropy import units as u -from astropy.io import ascii +from astropy.io import ascii as ascii_ from astropy.units import imperial from astropy.utils.data import get_pkg_data_filename import numpy as np @@ -78,8 +78,8 @@ alpha = 34.1632 * u.K / u.km # Reading layer parameters file -coesa_file = get_pkg_data_filename("data/coesa62.dat") -coesa62_data = ascii.read(coesa_file) +coesa_file = get_pkg_data_filename("../../core/earth/atmosphere/data/coesa62.dat") +coesa62_data = ascii_.read(coesa_file) b_levels = coesa62_data["b"].data zb_levels = coesa62_data["Zb [km]"].data * u.km hb_levels = coesa62_data["Hb [km]"].data * u.km diff --git a/src/hapsira/earth/atmosphere/coesa76.py b/src/hapsira/earth/atmosphere/coesa76.py index 465d5d435..efc30b2ed 100644 --- a/src/hapsira/earth/atmosphere/coesa76.py +++ b/src/hapsira/earth/atmosphere/coesa76.py @@ -44,58 +44,66 @@ """ from astropy import units as u -from astropy.io import ascii -from astropy.utils.data import get_pkg_data_filename import numpy as np from hapsira.earth.atmosphere.base import COESA -# Following constants come from original U.S Atmosphere 1962 paper so a pure -# model of this atmosphere can be implemented -R = 8314.32 * u.J / u.kmol / u.K -R_air = 287.053 * u.J / u.kg / u.K -k = 1.380622e-23 * u.J / u.K -Na = 6.022169e-26 / u.kmol -g0 = 9.80665 * u.m / u.s**2 -r0 = 6356.766 * u.km -M0 = 28.9644 * u.kg / u.kmol -P0 = 101325 * u.Pa -T0 = 288.15 * u.K -Tinf = 1000 * u.K -gamma = 1.4 -alpha = 34.1632 * u.K / u.km -beta = 1.458e-6 * (u.kg / u.s / u.m / (u.K) ** 0.5) -S = 110.4 * u.K +from hapsira.core.earth.atmosphere.coesa76 import ( + R, + R_air, + k, + Na, + g0, + r0, + M0, + P0, + T0, + Tinf, + gamma, + alpha, + beta, + S, + b_levels, + zb_levels, + hb_levels, + Tb_levels, + Lb_levels, + pb_levels, + z_coeff, + p_coeff, + rho_coeff, +) + +__all__ = [ + "COESA76", +] + +R = R * u.J / u.kmol / u.K +R_air = R_air * u.J / u.kg / u.K +k = k * u.J / u.K +Na = Na / u.kmol +g0 = g0 * u.m / u.s**2 +r0 = r0 * u.km +M0 = M0 * u.kg / u.kmol +P0 = P0 * u.Pa +T0 = T0 * u.K +Tinf = Tinf * u.K +alpha = alpha * u.K / u.km +beta = beta * (u.kg / u.s / u.m / (u.K) ** 0.5) +S = S * u.K # Reading layer parameters file -coesa76_data = ascii.read(get_pkg_data_filename("data/coesa76.dat")) -b_levels = coesa76_data["b"].data -zb_levels = coesa76_data["Zb [km]"].data * u.km -hb_levels = coesa76_data["Hb [km]"].data * u.km -Tb_levels = coesa76_data["Tb [K]"].data * u.K -Lb_levels = coesa76_data["Lb [K/km]"].data * u.K / u.km -pb_levels = coesa76_data["pb [mbar]"].data * u.mbar - -# Reading pressure and density coefficients files -p_data = ascii.read(get_pkg_data_filename("data/coesa76_p.dat")) -rho_data = ascii.read(get_pkg_data_filename("data/coesa76_rho.dat")) +b_levels = np.array(b_levels) +zb_levels = np.array(zb_levels) * u.km +hb_levels = np.array(hb_levels) * u.km +Tb_levels = np.array(Tb_levels) * u.K +Lb_levels = np.array(Lb_levels) * u.K / u.km +pb_levels = np.array(pb_levels) * u.mbar # Zip coefficients for each altitude -z_coeff = p_data["z [km]"].data * u.km -p_coeff = [ - p_data["A"].data, - p_data["B"].data, - p_data["C"].data, - p_data["D"].data, - p_data["E"].data, -] -rho_coeff = [ - rho_data["A"].data, - rho_data["B"].data, - rho_data["C"].data, - rho_data["D"].data, - rho_data["E"].data, -] +z_coeff = z_coeff * u.km +p_coeff = [np.array(entry) for entry in p_coeff] +rho_coeff = [np.array(entry) for entry in rho_coeff] class COESA76(COESA): diff --git a/tests/tests_earth/tests_atmosphere/test_coesa76.py b/tests/tests_earth/tests_atmosphere/test_coesa76.py index 774a70498..52bd2feaa 100644 --- a/tests/tests_earth/tests_atmosphere/test_coesa76.py +++ b/tests/tests_earth/tests_atmosphere/test_coesa76.py @@ -4,6 +4,7 @@ from hapsira.earth.atmosphere import COESA76 from hapsira.earth.atmosphere.coesa76 import p_coeff, rho_coeff +from numpy import float32 as f4 coesa76 = COESA76() @@ -28,18 +29,18 @@ def test_get_index_coesa76(): def test_coefficients_over_86km(): # Expected pressure coefficients expected_p = [ - 9.814674e-11, - -1.654439e-07, - 1.148115e-04, - -0.05431334, - -2.011365, + f4(9.814674e-11), + f4(-1.654439e-07), + f4(1.148115e-04), + f4(-0.05431334), + f4(-2.011365), ] expected_rho = [ - 1.140564e-10, - -2.130756e-07, - 1.570762e-04, - -0.07029296, - -12.89844, + f4(1.140564e-10), + f4(-2.130756e-07), + f4(1.570762e-04), + f4(-0.07029296), + f4(-12.89844), ] assert coesa76._get_coefficients_avobe_86(350 * u.km, p_coeff) == expected_p From 0ca0b76a46aada7794d9155869179711470e7f0a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 09:45:36 +0100 Subject: [PATCH 156/346] compiled coesa76 core functions, now works for cowell's method --- src/hapsira/core/earth/atmosphere/coesa76.py | 328 ++++++++++++++++-- .../tests_atmosphere/test_coesa76.py | 28 +- tests/tests_twobody/test_perturbations.py | 88 +++-- 3 files changed, 355 insertions(+), 89 deletions(-) diff --git a/src/hapsira/core/earth/atmosphere/coesa76.py b/src/hapsira/core/earth/atmosphere/coesa76.py index 7fc6e1b82..dc1d7acbe 100644 --- a/src/hapsira/core/earth/atmosphere/coesa76.py +++ b/src/hapsira/core/earth/atmosphere/coesa76.py @@ -1,7 +1,12 @@ +from math import exp + from astropy.io import ascii as ascii_ from astropy.utils.data import get_pkg_data_filename -from numpy import float32 as f4 +from numpy import float32 as float_, int64 as i8 + +from .coesa import check_altitude_hf +from ...jit import hjit __all__ = [ "R", @@ -27,51 +32,312 @@ "z_coeff", "p_coeff", "rho_coeff", + "pressure_hf", + "temperature_hf", + "density_hf", ] # Following constants come from original U.S Atmosphere 1962 paper so a pure # model of this atmosphere can be implemented -R = f4(8314.32) # u.J / u.kmol / u.K -R_air = f4(287.053) # u.J / u.kg / u.K -k = f4(1.380622e-23) # u.J / u.K -Na = f4(6.022169e-26) # 1 / u.kmol -g0 = f4(9.80665) # u.m / u.s**2 -r0 = f4(6356.766) # u.km -M0 = f4(28.9644) # u.kg / u.kmol -P0 = f4(101325) # u.Pa -T0 = f4(288.15) # u.K +R = float_(8314.32) # u.J / u.kmol / u.K +R_air = float_(287.053) # u.J / u.kg / u.K +k = float_(1.380622e-23) # u.J / u.K +Na = float_(6.022169e-26) # 1 / u.kmol +g0 = float_(9.80665) # u.m / u.s**2 +r0 = float_(6356.766) # u.km +M0 = float_(28.9644) # u.kg / u.kmol +P0 = float_(101325) # u.Pa +T0 = float_(288.15) # u.K Tinf = 1000 # u.K -gamma = f4(1.4) # one -alpha = f4(34.1632) # u.K / u.km -beta = f4(1.458e-6) # (u.kg / u.s / u.m / (u.K) ** 0.5) -S = f4(110.4) # u.K +gamma = float_(1.4) # one +alpha = float_(34.1632) # u.K / u.km +beta = float_(1.458e-6) # (u.kg / u.s / u.m / (u.K) ** 0.5) +S = float_(110.4) # u.K # Reading layer parameters file coesa76_data = ascii_.read(get_pkg_data_filename("data/coesa76.dat")) -b_levels = tuple(f4(number) for number in coesa76_data["b"].data) -zb_levels = tuple(f4(number) for number in coesa76_data["Zb [km]"].data) # u.km -hb_levels = tuple(f4(number) for number in coesa76_data["Hb [km]"].data) # u.km -Tb_levels = tuple(f4(number) for number in coesa76_data["Tb [K]"].data) # u.K -Lb_levels = tuple(f4(number) for number in coesa76_data["Lb [K/km]"].data) # u.K / u.km -pb_levels = tuple(f4(number) for number in coesa76_data["pb [mbar]"].data) # u.mbar +b_levels = tuple(i8(number) for number in coesa76_data["b"].data) +zb_levels = tuple(float_(number) for number in coesa76_data["Zb [km]"].data) # u.km +hb_levels = tuple(float_(number) for number in coesa76_data["Hb [km]"].data) # u.km +Tb_levels = tuple(float_(number) for number in coesa76_data["Tb [K]"].data) # u.K +Lb_levels = tuple( + float_(number) for number in coesa76_data["Lb [K/km]"].data +) # u.K / u.km +pb_levels = tuple(float_(number) for number in coesa76_data["pb [mbar]"].data) # u.mbar # Reading pressure and density coefficients files p_data = ascii_.read(get_pkg_data_filename("data/coesa76_p.dat")) rho_data = ascii_.read(get_pkg_data_filename("data/coesa76_rho.dat")) # Zip coefficients for each altitude -z_coeff = tuple(f4(number) for number in p_data["z [km]"].data) # u.km +z_coeff = tuple(i8(number) for number in p_data["z [km]"].data) # u.km p_coeff = ( - tuple(f4(number) for number in p_data["A"].data), - tuple(f4(number) for number in p_data["B"].data), - tuple(f4(number) for number in p_data["C"].data), - tuple(f4(number) for number in p_data["D"].data), - tuple(f4(number) for number in p_data["E"].data), + tuple(float_(number) for number in p_data["A"].data), + tuple(float_(number) for number in p_data["B"].data), + tuple(float_(number) for number in p_data["C"].data), + tuple(float_(number) for number in p_data["D"].data), + tuple(float_(number) for number in p_data["E"].data), ) rho_coeff = ( - tuple(f4(number) for number in rho_data["A"].data), - tuple(f4(number) for number in rho_data["B"].data), - tuple(f4(number) for number in rho_data["C"].data), - tuple(f4(number) for number in rho_data["D"].data), - tuple(f4(number) for number in rho_data["E"].data), + tuple(float_(number) for number in rho_data["A"].data), + tuple(float_(number) for number in rho_data["B"].data), + tuple(float_(number) for number in rho_data["C"].data), + tuple(float_(number) for number in rho_data["D"].data), + tuple(float_(number) for number in rho_data["E"].data), ) + + +@hjit("Tuple([f,f])(f,f,b1)") +def _check_altitude_hf(alt, r0, geometric): # geometric True by default + """Checks if altitude is inside valid range. + + Parameters + ---------- + alt : float + Altitude to be checked. + r0 : float + Attractor radius. + geometric : bool + If `True`, assumes geometric altitude kind. + + Returns + ------- + z: float + Geometric altitude. + h: float + Geopotential altitude. + + """ + + z, h = check_altitude_hf(alt, r0, geometric) + assert zb_levels[0] <= z <= zb_levels[-1] + + return z, h + + +@hjit("i8(f)") +def _get_index_zb_levels_hf(x): + """Finds element in list and returns index. + + Parameters + ---------- + x : float + Element to be searched. + + Returns + ------- + i: int + Index for the value. + + """ + for i, value in enumerate(zb_levels): + if i < len(zb_levels) and value < x: + continue + if x == value: + return i + return i - 1 + return 999 # HACK error ... ? + + +@hjit("i8(f)") +def _get_index_z_coeff_hf(x): + """Finds element in list and returns index. + + Parameters + ---------- + x : float + Element to be searched. + + Returns + ------- + i: int + Index for the value. + + """ + for i, value in enumerate(z_coeff): + if i < len(z_coeff) and value < x: + continue + if x == value: + return i + return i - 1 + return 999 # HACK error ... ? + + +@hjit("Tuple([f,f,f,f,f])(f)") +def _get_coefficients_avobe_86_p_coeff_hf(z): + """Returns corresponding coefficients for 4th order polynomial approximation. + + Parameters + ---------- + z : float + Geometric altitude + + Returns + ------- + coeffs: tuple[float,float,float,float,float] + List of corresponding coefficients + """ + # Get corresponding coefficients + i = _get_index_z_coeff_hf(z) + return p_coeff[0][i], p_coeff[1][i], p_coeff[2][i], p_coeff[3][i], p_coeff[4][i] + + +@hjit("Tuple([f,f,f,f,f])(f)") +def _get_coefficients_avobe_86_rho_coeff_hf(z): + """Returns corresponding coefficients for 4th order polynomial approximation. + + Parameters + ---------- + z : float + Geometric altitude + + Returns + ------- + coeffs: tuple[float,float,float,float,float] + List of corresponding coefficients + """ + # Get corresponding coefficients + i = _get_index_z_coeff_hf(z) + return ( + rho_coeff[0][i], + rho_coeff[1][i], + rho_coeff[2][i], + rho_coeff[3][i], + rho_coeff[4][i], + ) + + +@hjit("f(f,b1)") +def temperature_hf(alt, geometric): # geometric True by default + """Solves for temperature at given altitude. + + Parameters + ---------- + alt : float + Geometric/Geopotential altitude. + geometric : bool + If `True`, assumes geometric altitude kind. + + Returns + ------- + T: float + Kinetic temeperature. + """ + # Test if altitude is inside valid range + z, h = _check_altitude_hf(alt, r0, geometric) + + # Get base parameters + i = _get_index_zb_levels_hf(z) + Tb = Tb_levels[i] + Lb = Lb_levels[i] + hb = hb_levels[i] + + # Apply different equations + if z < zb_levels[7]: + # Below 86km + # TODO: Apply air mean molecular weight ratio factor + Tm = Tb + Lb * (h - hb) + T = Tm + elif zb_levels[7] <= z and z < zb_levels[8]: + # [86km, 91km) + T = 186.87 + elif zb_levels[8] <= z and z < zb_levels[9]: + # [91km, 110km] + Tc = 263.1905 + A = -76.3232 + a = -19.9429 + T = Tc + A * (1 - ((z - zb_levels[8]) / a) ** 2) ** 0.5 + elif zb_levels[9] <= z and z < zb_levels[10]: + # [110km, 120km] + T = 240 + Lb * (z - zb_levels[9]) + else: + T10 = 360.0 + _gamma = Lb_levels[9] / (Tinf - T10) + epsilon = (z - zb_levels[10]) * (r0 + zb_levels[10]) / (r0 + z) + T = Tinf - (Tinf - T10) * exp(-_gamma * epsilon) + + return T + + +@hjit("f(f,b1)") +def pressure_hf(alt, geometric): # geometric True by default + """Solves pressure at given altitude. + + Parameters + ---------- + alt : float + Geometric/Geopotential altitude. + geometric : bool + If `True`, assumes geometric altitude kind. + + Returns + ------- + p: float + Pressure at given altitude. + """ + # Test if altitude is inside valid range + z, h = _check_altitude_hf(alt, r0, geometric) + + # Obtain gravity magnitude + # Get base parameters + i = _get_index_zb_levels_hf(z) + Tb = Tb_levels[i] + Lb = Lb_levels[i] + hb = hb_levels[i] + pb = pb_levels[i] + + # If above 86[km] usual formulation is applied + if z < 86: + if Lb == 0.0: + p = pb * exp(-alpha * (h - hb) / Tb) * 100 # HACK 100 ... SI-prefix change? + else: + T = temperature_hf(z, geometric) + p = pb * (Tb / T) ** (alpha / Lb) * 100 # HACK 100 ... SI-prefix change? + else: + # TODO: equation (33c) should be applied instead of using coefficients + + # A 4th order polynomial is used to approximate pressure. This was + # directly taken from: http://www.braeunig.us/space/atmmodel.htm + A, B, C, D, E = _get_coefficients_avobe_86_p_coeff_hf(z) + + # Solve the polynomial + p = exp(A * z**4 + B * z**3 + C * z**2 + D * z + E) + + return p + + +@hjit("f(f,b1)") +def density_hf(alt, geometric): # geometric True by default + """Solves density at given height. + + Parameters + ---------- + alt : float + Geometric/Geopotential height. + geometric : bool + If `True`, assumes that `alt` argument is geometric kind. + + Returns + ------- + rho: float + Density at given height. + """ + # Test if altitude is inside valid range + z, _ = _check_altitude_hf(alt, r0, geometric) + + # Solve temperature and pressure + if z <= 86: + T = temperature_hf(z, geometric) + p = pressure_hf(z, geometric) + rho = p / R_air / T + else: + # TODO: equation (42) should be applied instead of using coefficients + + # A 4th order polynomial is used to approximate pressure. This was + # directly taken from: http://www.braeunig.us/space/atmmodel.htm + A, B, C, D, E = _get_coefficients_avobe_86_rho_coeff_hf(z) + + # Solve the polynomial + rho = exp(A * z**4 + B * z**3 + C * z**2 + D * z + E) + + return rho diff --git a/tests/tests_earth/tests_atmosphere/test_coesa76.py b/tests/tests_earth/tests_atmosphere/test_coesa76.py index 52bd2feaa..962da527a 100644 --- a/tests/tests_earth/tests_atmosphere/test_coesa76.py +++ b/tests/tests_earth/tests_atmosphere/test_coesa76.py @@ -1,10 +1,10 @@ from astropy import units as u from astropy.tests.helper import assert_quantity_allclose +from numpy.testing import assert_allclose import pytest from hapsira.earth.atmosphere import COESA76 from hapsira.earth.atmosphere.coesa76 import p_coeff, rho_coeff -from numpy import float32 as f4 coesa76 = COESA76() @@ -29,22 +29,24 @@ def test_get_index_coesa76(): def test_coefficients_over_86km(): # Expected pressure coefficients expected_p = [ - f4(9.814674e-11), - f4(-1.654439e-07), - f4(1.148115e-04), - f4(-0.05431334), - f4(-2.011365), + 9.814674e-11, + -1.654439e-07, + 1.148115e-04, + -0.05431334, + -2.011365, ] expected_rho = [ - f4(1.140564e-10), - f4(-2.130756e-07), - f4(1.570762e-04), - f4(-0.07029296), - f4(-12.89844), + 1.140564e-10, + -2.130756e-07, + 1.570762e-04, + -0.07029296, + -12.89844, ] - assert coesa76._get_coefficients_avobe_86(350 * u.km, p_coeff) == expected_p - assert coesa76._get_coefficients_avobe_86(350 * u.km, rho_coeff) == expected_rho + assert_allclose(coesa76._get_coefficients_avobe_86(350 * u.km, p_coeff), expected_p) + assert_allclose( + coesa76._get_coefficients_avobe_86(350 * u.km, rho_coeff), expected_rho + ) # SOLUTIONS DIRECTLY TAKEN FROM COESA76 REPORT diff --git a/tests/tests_twobody/test_perturbations.py b/tests/tests_twobody/test_perturbations.py index e89b048a8..0cd525f45 100644 --- a/tests/tests_twobody/test_perturbations.py +++ b/tests/tests_twobody/test_perturbations.py @@ -8,20 +8,20 @@ from hapsira.bodies import Earth, Moon, Sun from hapsira.constants import H0_earth, Wdivc_sun, rho0_earth +from hapsira.core.earth.atmosphere.coesa76 import density_hf as coesa76_density_hf from hapsira.core.elements import rv2coe_gf, RV2COE_TOL from hapsira.core.jit import hjit, djit from hapsira.core.math.linalg import add_VV_hf, mul_Vs_hf, norm_hf from hapsira.core.perturbations import ( # pylint: disable=E1120,E1136 J2_perturbation_hf, J3_perturbation_hf, - # atmospheric_drag_hf, # TODO reactivate test + atmospheric_drag_hf, atmospheric_drag_exponential_hf, radiation_pressure_hf, third_body_hf, ) from hapsira.core.propagation.base import func_twobody_hf -# from hapsira.earth.atmosphere import COESA76 # TODO reactivate test from hapsira.ephem import build_ephem_interpolant from hapsira.twobody import Orbit from hapsira.twobody.events import LithobrakeEvent @@ -361,60 +361,58 @@ def f_hf(t0, rr, vv, k): assert lithobrake_event.last_t == tofs[-1] -# @pytest.mark.slow -# def test_atmospheric_demise_coesa76(): -# # Test an orbital decay that hits Earth. No analytic solution. -# R = Earth.R.to(u.km).value - -# orbit = Orbit.circular(Earth, 250 * u.km) -# t_decay = 7.17 * u.d +@pytest.mark.slow +def test_atmospheric_demise_coesa76(): + # Test an orbital decay that hits Earth. No analytic solution. + R = Earth.R.to(u.km).value -# # Parameters of a body -# C_D = 2.2 # Dimensionless (any value would do) -# A_over_m = ((np.pi / 4.0) * (u.m**2) / (100 * u.kg)).to_value( -# u.km**2 / u.kg -# ) # km^2/kg + orbit = Orbit.circular(Earth, 250 * u.km) + t_decay = 7.17 * u.d -# tofs = [365] * u.d + # Parameters of a body + C_D = 2.2 # Dimensionless (any value would do) + A_over_m = ((np.pi / 4.0) * (u.m**2) / (100 * u.kg)).to_value( + u.km**2 / u.kg + ) # km^2/kg -# lithobrake_event = LithobrakeEvent(R) -# events = [lithobrake_event] + tofs = [365] * u.d -# coesa76 = COESA76() + lithobrake_event = LithobrakeEvent(R) + events = [lithobrake_event] -# @djit -# def f_hf(t0, rr, vv, k): -# du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + @djit + def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) -# # Avoid undershooting H below attractor radius R -# H = norm_hf(rr) -# if H < R: -# H = R + # Avoid undershooting H below attractor radius R + H = norm_hf(rr) + if H < R: + H = R -# rho = coesa76.density( -# (H - R) * u.km -# ).to_value(u.kg / u.km**3) # TODO jit'ed ... move to core?!? + rho = ( + coesa76_density_hf(H - R, True) * 1e9 + ) # HACK convert from kg/m**3 to kg/km**3 -# du_ad = atmospheric_drag_hf( -# t0, -# rr, -# vv, -# k, -# C_D, -# A_over_m, -# rho, -# ) -# return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) + du_ad = atmospheric_drag_hf( + t0, + rr, + vv, + k, + C_D, + A_over_m, + rho, + ) + return du_kep_rr, add_VV_hf(du_kep_vv, du_ad) -# method = CowellPropagator(events=events, f=f) -# rr, _ = method.propagate_many( -# orbit._state, -# tofs, -# ) + method = CowellPropagator(events=events, f=f_hf) + rr, _ = method.propagate_many( + orbit._state, + tofs, + ) -# assert_quantity_allclose(norm(rr[0].to(u.km).value), R, atol=1) # Below 1km + assert_quantity_allclose(norm(rr[0].to(u.km).value), R, atol=1) # Below 1km -# assert_quantity_allclose(lithobrake_event.last_t, t_decay, rtol=1e-2) + assert_quantity_allclose(lithobrake_event.last_t, t_decay, rtol=1e-2) @pytest.mark.slow From 486c556b2c983eabfabb3e43539df4e59d945b3a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 10:52:20 +0100 Subject: [PATCH 157/346] use compiled coesa76 funcs in modell class --- src/hapsira/core/earth/atmosphere/coesa76.py | 35 +++++++- src/hapsira/earth/atmosphere/coesa76.py | 91 ++------------------ 2 files changed, 40 insertions(+), 86 deletions(-) diff --git a/src/hapsira/core/earth/atmosphere/coesa76.py b/src/hapsira/core/earth/atmosphere/coesa76.py index dc1d7acbe..4c1c9752d 100644 --- a/src/hapsira/core/earth/atmosphere/coesa76.py +++ b/src/hapsira/core/earth/atmosphere/coesa76.py @@ -6,7 +6,7 @@ from numpy import float32 as float_, int64 as i8 from .coesa import check_altitude_hf -from ...jit import hjit +from ...jit import hjit, vjit __all__ = [ "R", @@ -32,9 +32,13 @@ "z_coeff", "p_coeff", "rho_coeff", + "COESA76_GEOMETRIC", "pressure_hf", + "pressure_vf", "temperature_hf", + "temperature_vf", "density_hf", + "density_vf", ] # Following constants come from original U.S Atmosphere 1962 paper so a pure @@ -86,6 +90,8 @@ tuple(float_(number) for number in rho_data["E"].data), ) +COESA76_GEOMETRIC = True + @hjit("Tuple([f,f])(f,f,b1)") def _check_altitude_hf(alt, r0, geometric): # geometric True by default @@ -259,6 +265,15 @@ def temperature_hf(alt, geometric): # geometric True by default return T +@vjit("f(f,b1)") +def temperature_vf(alt, geometric): # geometric True by default + """ + Vectorized temperature + """ + + return temperature_hf(alt, geometric) + + @hjit("f(f,b1)") def pressure_hf(alt, geometric): # geometric True by default """Solves pressure at given altitude. @@ -306,6 +321,15 @@ def pressure_hf(alt, geometric): # geometric True by default return p +@vjit("f(f,b1)") +def pressure_vf(alt, geometric): # geometric True by default + """ + Vectorized pressure + """ + + return pressure_hf(alt, geometric) + + @hjit("f(f,b1)") def density_hf(alt, geometric): # geometric True by default """Solves density at given height. @@ -341,3 +365,12 @@ def density_hf(alt, geometric): # geometric True by default rho = exp(A * z**4 + B * z**3 + C * z**2 + D * z + E) return rho + + +@vjit("f(f,b1)") +def density_vf(alt, geometric): # geometric True by default + """ + Vectorized density + """ + + return density_hf(alt, geometric) diff --git a/src/hapsira/earth/atmosphere/coesa76.py b/src/hapsira/earth/atmosphere/coesa76.py index efc30b2ed..95cb2bc90 100644 --- a/src/hapsira/earth/atmosphere/coesa76.py +++ b/src/hapsira/earth/atmosphere/coesa76.py @@ -72,6 +72,9 @@ z_coeff, p_coeff, rho_coeff, + pressure_vf, + density_vf, + temperature_vf, ) __all__ = [ @@ -153,40 +156,8 @@ def temperature(self, alt, geometric=True): T: ~astropy.units.Quantity Kinetic temeperature. """ - # Test if altitude is inside valid range - z, h = self._check_altitude(alt, r0, geometric=geometric) - # Get base parameters - i = self._get_index(z, self.zb_levels) - Tb = self.Tb_levels[i] - Lb = self.Lb_levels[i] - hb = self.hb_levels[i] - - # Apply different equations - if z < self.zb_levels[7]: - # Below 86km - # TODO: Apply air mean molecular weight ratio factor - Tm = Tb + Lb * (h - hb) - T = Tm - elif self.zb_levels[7] <= z and z < self.zb_levels[8]: - # [86km, 91km) - T = 186.87 * u.K - elif self.zb_levels[8] <= z and z < self.zb_levels[9]: - # [91km, 110km] - Tc = 263.1905 * u.K - A = -76.3232 * u.K - a = -19.9429 * u.km - T = Tc + A * (1 - ((z - self.zb_levels[8]) / a) ** 2) ** 0.5 - elif self.zb_levels[9] <= z and z < self.zb_levels[10]: - # [110km, 120km] - T = 240 * u.K + Lb * (z - self.zb_levels[9]) - else: - T10 = 360.0 * u.K - _gamma = self.Lb_levels[9] / (Tinf - T10) - epsilon = (z - self.zb_levels[10]) * (r0 + self.zb_levels[10]) / (r0 + z) - T = Tinf - (Tinf - T10) * np.exp(-_gamma * epsilon) - - return T.to(u.K) + return temperature_vf(alt.to_value(u.km), geometric) * u.K def pressure(self, alt, geometric=True): """Solves pressure at given altitude. @@ -203,36 +174,8 @@ def pressure(self, alt, geometric=True): p: ~astropy.units.Quantity Pressure at given altitude. """ - # Test if altitude is inside valid range - z, h = self._check_altitude(alt, r0, geometric=geometric) - # Obtain gravity magnitude - # Get base parameters - i = self._get_index(z, self.zb_levels) - Tb = self.Tb_levels[i] - Lb = self.Lb_levels[i] - hb = self.hb_levels[i] - pb = self.pb_levels[i] - - # If above 86[km] usual formulation is applied - if z < 86 * u.km: - if Lb == 0.0 * u.K / u.km: - p = pb * np.exp(-alpha * (h - hb) / Tb) - else: - T = self.temperature(z) - p = pb * (Tb / T) ** (alpha / Lb) - else: - # TODO: equation (33c) should be applied instead of using coefficients - - # A 4th order polynomial is used to approximate pressure. This was - # directly taken from: http://www.braeunig.us/space/atmmodel.htm - A, B, C, D, E = self._get_coefficients_avobe_86(z, p_coeff) - - # Solve the polynomial - z = z.to_value(u.km) - p = np.exp(A * z**4 + B * z**3 + C * z**2 + D * z + E) * u.Pa - - return p.to(u.Pa) + return pressure_vf(alt.to_value(u.km), geometric) * u.Pa def density(self, alt, geometric=True): """Solves density at given height. @@ -249,30 +192,8 @@ def density(self, alt, geometric=True): rho: ~astropy.units.Quantity Density at given height. """ - # Test if altitude is inside valid range - z, h = self._check_altitude(alt, r0, geometric=geometric) - - # Solve temperature and pressure - if z <= 86 * u.km: - T = self.temperature(z) - p = self.pressure(z) - rho = p / R_air / T - else: - # TODO: equation (42) should be applied instead of using coefficients - - # A 4th order polynomial is used to approximate pressure. This was - # directly taken from: http://www.braeunig.us/space/atmmodel.htm - A, B, C, D, E = self._get_coefficients_avobe_86(z, rho_coeff) - - # Solve the polynomial - z = z.to_value(u.km) - rho = ( - np.exp(A * z**4 + B * z**3 + C * z**2 + D * z + E) - * u.kg - / u.m**3 - ) - return rho.to(u.kg / u.m**3) + return density_vf(alt.to_value(u.km), geometric) * u.kg / u.m**3 def properties(self, alt, geometric=True): """Solves temperature, pressure, density at given height. From af48adbbb83dabaf58608bd2b32dcd8fcb634fed Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 11:02:04 +0100 Subject: [PATCH 158/346] rm notations --- src/hapsira/core/math/ivp/_rkstep.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index 178745b5b..382bdacf6 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -1,5 +1,3 @@ -from typing import Callable, Tuple - from ._dop853_coefficients import A as _A, B as _B, C as _C from ..linalg import add_VV_hf from ...jit import hjit, DSIG @@ -46,16 +44,7 @@ @hjit(f"Tuple([V,V,V,V,{_KSIG:s}])(F({DSIG:s}),f,V,V,V,V,f,f)") -def rk_step_hf( - fun: Callable, - t: float, - rr: tuple[float, float, float], - vv: tuple[float, float, float], - fr: tuple[float, float, float], - fv: tuple[float, float, float], - h: float, - argk: float, -) -> Tuple[Tuple, Tuple, Tuple, Tuple, Tuple]: +def rk_step_hf(fun, t, rr, vv, fr, fv, h, argk): """Perform a single Runge-Kutta step. This function computes a prediction of an explicit Runge-Kutta method and From b7b178f7eb7b06e6b425bf95ae4801d6b38edc53 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 11:50:39 +0100 Subject: [PATCH 159/346] rk initial step --- src/hapsira/core/math/ivp/_rk.py | 134 +++++++++++++------------------ src/hapsira/core/math/linalg.py | 6 ++ 2 files changed, 64 insertions(+), 76 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index ddbb9de5a..de4cdad1b 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -1,13 +1,13 @@ -from typing import Callable, Tuple -from warnings import warn +from math import sqrt +from typing import Callable import numpy as np -from numba import jit from . import _dop853_coefficients as dop853_coefficients from ._rkstep import rk_step_hf, N_RV, N_STAGES -from ...jit import array_to_V_hf +from ...jit import array_to_V_hf, hjit, DSIG +from ...math.linalg import add_VV_hf, div_VV_hf, mul_Vs_hf, sub_VV_hf __all__ = [ "EPS", @@ -28,24 +28,15 @@ ERROR_ESTIMATOR_ORDER = 7 -@jit(nopython=False) -def norm(x: np.ndarray) -> float: - """Compute RMS norm.""" - return np.linalg.norm(x) / x.size**0.5 - - -@jit(nopython=False) -def select_initial_step( - fun: Callable, - t0: float, - y0: np.ndarray, - argk: float, - f0: np.ndarray, - direction: float, - order: float, - rtol: float, - atol: float, -) -> float: +@hjit("f(V,V)") +def _norm_VV_hf(x, y): + return sqrt(x[0] ** 2 + x[1] ** 2 + x[2] ** 2 + y[0] ** 2 + y[1] ** 2 + y[2] ** 2) + + +@hjit(f"f(F({DSIG:s}),f,V,V,f,V,V,f,f,f,f)") +def _select_initial_step_hf( + fun, t0, rr, vv, argk, fr, fv, direction, order, rtol, atol +): """Empirically select a good initial step. The algorithm is described in [1]_. @@ -80,26 +71,44 @@ def select_initial_step( .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential Equations I: Nonstiff Problems", Sec. II.4. """ - if y0.size == 0: - return np.inf - scale = atol + np.abs(y0) * rtol - d0 = norm(y0 / scale) - d1 = norm(f0 / scale) + scale_r = ( + atol + abs(rr[0]) * rtol, + atol + abs(rr[1]) * rtol, + atol + abs(rr[2]) * rtol, + ) + scale_v = ( + atol + abs(vv[0]) * rtol, + atol + abs(vv[1]) * rtol, + atol + abs(vv[2]) * rtol, + ) + + factor = 1 / sqrt(6) + d0 = _norm_VV_hf(div_VV_hf(rr, scale_r), div_VV_hf(vv, scale_v)) * factor + d1 = _norm_VV_hf(div_VV_hf(fr, scale_r), div_VV_hf(fv, scale_v)) * factor + if d0 < 1e-5 or d1 < 1e-5: h0 = 1e-6 else: h0 = 0.01 * d0 / d1 - y1 = y0 + h0 * direction * f0 - rr, vv = fun( + yr1 = add_VV_hf(rr, mul_Vs_hf(fr, h0 * direction)) + yv1 = add_VV_hf(vv, mul_Vs_hf(fv, h0 * direction)) + + fr1, fv1 = fun( t0 + h0 * direction, - array_to_V_hf(y1[:3]), - array_to_V_hf(y1[3:]), + yr1, + yv1, argk, - ) # TODO call into hf - f1 = np.array([*rr, *vv]) - d2 = norm((f1 - f0) / scale) / h0 + ) + + d2 = ( + _norm_VV_hf( + div_VV_hf(sub_VV_hf(fr1, fr), scale_r), + div_VV_hf(sub_VV_hf(fv1, fv), scale_v), + ) + / h0 + ) if d1 <= 1e-15 and d2 <= 1e-15: h1 = max(1e-6, h0 * 1e-3) @@ -109,35 +118,6 @@ def select_initial_step( return min(100 * h0, h1) -@jit(nopython=False) -def validate_max_step(max_step: float) -> float: - """Assert that max_Step is valid and return it.""" - if max_step <= 0: - raise ValueError("`max_step` must be positive.") - return max_step - - -@jit(nopython=False) -def validate_tol(rtol: float, atol: float) -> Tuple[float, float]: - """Validate tolerance values.""" - - if np.any(rtol < 100 * EPS): - warn( - "At least one element of `rtol` is too small. " - f"Setting `rtol = np.maximum(rtol, {100 * EPS})`." - ) - rtol = np.maximum(rtol, 100 * EPS) - - atol = np.asarray(atol) - if atol.ndim > 0 and atol.shape != (N_RV,): - raise ValueError("`atol` has wrong shape.") - - if np.any(atol < 0): - raise ValueError("`atol` must be positive.") - - return rtol, atol - - class Dop853DenseOutput: """local interpolant over step made by an ODE solver. @@ -312,8 +292,15 @@ def __init__( self.nlu = 0 self.y_old = None - self.max_step = validate_max_step(max_step) - self.rtol, self.atol = validate_tol(rtol, atol) + + assert max_step > 0 + self.max_step = max_step + + if rtol < 100 * EPS: + rtol = 100 * EPS + assert atol >= 0 + self.rtol, self.atol = rtol, atol + rr, vv = self.fun( self.t, array_to_V_hf(self.y[:3]), @@ -321,30 +308,25 @@ def __init__( self.argk, ) # TODO call into hf self.f = np.array([*rr, *vv]) - self.h_abs = select_initial_step( + self.h_abs = _select_initial_step_hf( self.fun, self.t, - self.y, + array_to_V_hf(self.y[:3]), + array_to_V_hf(self.y[3:]), self.argk, - self.f, + array_to_V_hf(self.f[:3]), + array_to_V_hf(self.f[3:]), self.direction, ERROR_ESTIMATOR_ORDER, self.rtol, self.atol, - ) + ) # TODO call into hf self.error_exponent = -1 / (ERROR_ESTIMATOR_ORDER + 1) self.h_previous = None self.K_extended = np.empty((N_STAGES_EXTENDED, N_RV), dtype=self.y.dtype) self.K = self.K_extended[: N_STAGES + 1] - @property - def step_size(self): - if self.t_old is None: - return None - else: - return np.abs(self.t - self.t_old) - def step(self): """Perform one integration step. diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index f856ba1ed..7c6ea5092 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -6,6 +6,7 @@ "add_VV_hf", "cross_VV_hf", "div_Vs_hf", + "div_VV_hf", "matmul_MM_hf", "matmul_VM_hf", "matmul_VV_hf", @@ -43,6 +44,11 @@ def div_ss_hf(a, b): return a / b +@hjit("V(V,V)") +def div_VV_hf(x, y): + return x[0] / y[0], x[1] / y[1], x[2] / y[2] + + @hjit("V(V,f)") def div_Vs_hf(v, s): return v[0] / s, v[1] / s, v[2] / s From d62994db1f6a073be179f9d3ccb2f1bf2022a530 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 12:01:16 +0100 Subject: [PATCH 160/346] rm comment --- src/hapsira/core/math/ivp/_rk.py | 35 -------------------------------- 1 file changed, 35 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index de4cdad1b..f82d78786 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -37,41 +37,6 @@ def _norm_VV_hf(x, y): def _select_initial_step_hf( fun, t0, rr, vv, argk, fr, fv, direction, order, rtol, atol ): - """Empirically select a good initial step. - - The algorithm is described in [1]_. - - Parameters - ---------- - fun : callable - Right-hand side of the system. - t0 : float - Initial value of the independent variable. - y0 : ndarray, shape (n,) - Initial value of the dependent variable. - f0 : ndarray, shape (n,) - Initial value of the derivative, i.e., ``fun(t0, y0)``. - direction : float - Integration direction. - order : float - Error estimator order. It means that the error controlled by the - algorithm is proportional to ``step_size ** (order + 1)`. - rtol : float - Desired relative tolerance. - atol : float - Desired absolute tolerance. - - Returns - ------- - h_abs : float - Absolute value of the suggested initial step. - - References - ---------- - .. [1] E. Hairer, S. P. Norsett G. Wanner, "Solving Ordinary Differential - Equations I: Nonstiff Problems", Sec. II.4. - """ - scale_r = ( atol + abs(rr[0]) * rtol, atol + abs(rr[1]) * rtol, From 6b7af5d3d97bdf2d81e88b49e578854ae5fd4cf3 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 12:09:48 +0100 Subject: [PATCH 161/346] consolidate --- src/hapsira/core/math/ivp/_rk.py | 51 +++++++++++++++----------------- 1 file changed, 24 insertions(+), 27 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index f82d78786..a19455b19 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -336,7 +336,30 @@ def dense_output(self): assert self.t != self.t_old - return self._dense_output_impl() + K = self.K_extended + h = self.h_previous + for s, (a, c) in enumerate(zip(self.A_EXTRA, self.C_EXTRA), start=N_STAGES + 1): + dy = np.dot(K[:s].T, a[:s]) * h + y_ = self.y_old + dy + rr, vv = self.fun( + self.t_old + c * h, + array_to_V_hf(y_[:3]), + array_to_V_hf(y_[3:]), + self.argk, + ) # TODO call into hf + K[s] = np.array([*rr, *vv]) + + F = np.empty((INTERPOLATOR_POWER, N_RV), dtype=self.y_old.dtype) + + f_old = K[0] + delta_y = self.y - self.y_old + + F[0] = delta_y + F[1] = h * f_old - delta_y + F[2] = 2 * delta_y - h * (self.f + f_old) + F[3:] = h * np.dot(self.D, K) + + return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) def _estimate_error_norm(self, K, h, scale): err5 = np.dot(K.T, self.E5) / scale @@ -424,29 +447,3 @@ def _step_impl(self): self.f = f_new return True, None - - def _dense_output_impl(self): - K = self.K_extended - h = self.h_previous - for s, (a, c) in enumerate(zip(self.A_EXTRA, self.C_EXTRA), start=N_STAGES + 1): - dy = np.dot(K[:s].T, a[:s]) * h - y_ = self.y_old + dy - rr, vv = self.fun( - self.t_old + c * h, - array_to_V_hf(y_[:3]), - array_to_V_hf(y_[3:]), - self.argk, - ) # TODO call into hf - K[s] = np.array([*rr, *vv]) - - F = np.empty((INTERPOLATOR_POWER, N_RV), dtype=self.y_old.dtype) - - f_old = K[0] - delta_y = self.y - self.y_old - - F[0] = delta_y - F[1] = h * f_old - delta_y - F[2] = 2 * delta_y - h * (self.f + f_old) - F[3:] = h * np.dot(self.D, K) - - return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) From f0fe33c9fee86750b71b2ae70fa99bc0ef810b0b Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 12:27:35 +0100 Subject: [PATCH 162/346] rm method --- src/hapsira/core/math/ivp/_rk.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index a19455b19..8fac4e72c 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -117,9 +117,7 @@ def __call__(self, t): """ t = np.asarray(t) assert not t.ndim > 1 - return self._call_impl(t) - def _call_impl(self, t): x = (t - self.t_old) / self.h if t.ndim == 0: From 1b8a4f5c0aa9b0af82302fa0b945ef56f20d94b2 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 12:37:04 +0100 Subject: [PATCH 163/346] prep for compiled estimator --- src/hapsira/core/math/ivp/_rk.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 8fac4e72c..e743265b8 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -3,7 +3,7 @@ import numpy as np -from . import _dop853_coefficients as dop853_coefficients +from ._dop853_coefficients import E3, E5, A as _A, C as _C, D as _D from ._rkstep import rk_step_hf, N_RV, N_STAGES from ...jit import array_to_V_hf, hjit, DSIG @@ -219,12 +219,9 @@ class DOP853: TOO_SMALL_STEP = "Required step size is less than spacing between numbers." - E3 = dop853_coefficients.E3 - E5 = dop853_coefficients.E5 - D = dop853_coefficients.D - - A_EXTRA = dop853_coefficients.A[N_STAGES + 1 :] - C_EXTRA = dop853_coefficients.C[N_STAGES + 1 :] + A_EXTRA = _A[N_STAGES + 1 :] + C_EXTRA = _C[N_STAGES + 1 :] + D = _D def __init__( self, @@ -360,13 +357,20 @@ def dense_output(self): return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) def _estimate_error_norm(self, K, h, scale): - err5 = np.dot(K.T, self.E5) / scale - err3 = np.dot(K.T, self.E3) / scale + assert K.shape == (N_STAGES + 1, N_RV) + assert E3.shape == (N_STAGES + 1,) + assert E5.shape == (N_STAGES + 1,) + assert scale.shape == (N_RV,) + + err5 = np.dot(K.T, E5) / scale + err3 = np.dot(K.T, E3) / scale err5_norm_2 = np.linalg.norm(err5) ** 2 err3_norm_2 = np.linalg.norm(err3) ** 2 + if err5_norm_2 == 0 and err3_norm_2 == 0: return 0.0 denom = err5_norm_2 + 0.01 * err3_norm_2 + return np.abs(h) * err5_norm_2 / np.sqrt(denom * len(scale)) def _step_impl(self): From ed05028838137d365b400eef0abb6f8ffd4e0de5 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 14:10:52 +0100 Subject: [PATCH 164/346] jit error --- src/hapsira/core/math/ivp/_rk.py | 28 +-- src/hapsira/core/math/ivp/_rkerror.py | 249 ++++++++++++++++++++++++++ src/hapsira/core/math/ivp/_rkstep.py | 5 +- 3 files changed, 261 insertions(+), 21 deletions(-) create mode 100644 src/hapsira/core/math/ivp/_rkerror.py diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index e743265b8..c032ccaba 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -3,8 +3,9 @@ import numpy as np -from ._dop853_coefficients import E3, E5, A as _A, C as _C, D as _D +from ._dop853_coefficients import A as _A, C as _C, D as _D from ._rkstep import rk_step_hf, N_RV, N_STAGES +from ._rkerror import estimate_error_norm_hf from ...jit import array_to_V_hf, hjit, DSIG from ...math.linalg import add_VV_hf, div_VV_hf, mul_Vs_hf, sub_VV_hf @@ -356,23 +357,6 @@ def dense_output(self): return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) - def _estimate_error_norm(self, K, h, scale): - assert K.shape == (N_STAGES + 1, N_RV) - assert E3.shape == (N_STAGES + 1,) - assert E5.shape == (N_STAGES + 1,) - assert scale.shape == (N_RV,) - - err5 = np.dot(K.T, E5) / scale - err3 = np.dot(K.T, E3) / scale - err5_norm_2 = np.linalg.norm(err5) ** 2 - err3_norm_2 = np.linalg.norm(err3) ** 2 - - if err5_norm_2 == 0 and err3_norm_2 == 0: - return 0.0 - denom = err5_norm_2 + 0.01 * err3_norm_2 - - return np.abs(h) * err5_norm_2 / np.sqrt(denom * len(scale)) - def _step_impl(self): t = self.t y = self.y @@ -421,7 +405,13 @@ def _step_impl(self): self.K[: N_STAGES + 1, :N_RV] = np.array([K_new]) scale = atol + np.maximum(np.abs(y), np.abs(y_new)) * rtol - error_norm = self._estimate_error_norm(self.K, h, scale) + assert scale.shape == (N_RV,) + error_norm = estimate_error_norm_hf( + K_new, + h, + array_to_V_hf(scale[:3]), + array_to_V_hf(scale[3:]), + ) # TODO call into hf if error_norm < 1: if error_norm == 0: diff --git a/src/hapsira/core/math/ivp/_rkerror.py b/src/hapsira/core/math/ivp/_rkerror.py new file mode 100644 index 000000000..7a5a0a1cb --- /dev/null +++ b/src/hapsira/core/math/ivp/_rkerror.py @@ -0,0 +1,249 @@ +from math import sqrt + +from ._dop853_coefficients import E3 as _E3, E5 as _E5 +from ._rkstep import N_RV, KSIG +from ...jit import hjit +from ....settings import settings + +if settings["PRECISION"].value == "f8": + from numpy import float64 as float_ +elif settings["PRECISION"].value == "f4": + from numpy import float32 as float_ +elif settings["PRECISION"].value == "f2": + from numpy import float16 as float_ +else: + raise ValueError("unsupported precision") + + +__all__ = [ + "estimate_error_norm_hf", +] + + +E3 = tuple(float_(number) for number in _E3) # N_STAGES + 1 +E5 = tuple(float_(number) for number in _E5) # N_STAGES + 1 + + +@hjit(f"f({KSIG:s},f,V,V)") +def estimate_error_norm_hf(K, h, scale_r, scale_v): + K00, K01, K02, K03, K04, K05, K06, K07, K08, K09, K10, K11, K12 = K + + err3 = ( + ( + K00[0] * E3[0] + + K01[0] * E3[1] + + K02[0] * E3[2] + + K03[0] * E3[3] + + K04[0] * E3[4] + + K05[0] * E3[5] + + K06[0] * E3[6] + + K07[0] * E3[7] + + K08[0] * E3[8] + + K09[0] * E3[9] + + K10[0] * E3[10] + + K11[0] * E3[11] + + K12[0] * E3[12] + ) + / scale_r[0], + ( + K00[1] * E3[0] + + K01[1] * E3[1] + + K02[1] * E3[2] + + K03[1] * E3[3] + + K04[1] * E3[4] + + K05[1] * E3[5] + + K06[1] * E3[6] + + K07[1] * E3[7] + + K08[1] * E3[8] + + K09[1] * E3[9] + + K10[1] * E3[10] + + K11[1] * E3[11] + + K12[1] * E3[12] + ) + / scale_r[1], + ( + K00[2] * E3[0] + + K01[2] * E3[1] + + K02[2] * E3[2] + + K03[2] * E3[3] + + K04[2] * E3[4] + + K05[2] * E3[5] + + K06[2] * E3[6] + + K07[2] * E3[7] + + K08[2] * E3[8] + + K09[2] * E3[9] + + K10[2] * E3[10] + + K11[2] * E3[11] + + K12[2] * E3[12] + ) + / scale_r[2], + ( + K00[3] * E3[0] + + K01[3] * E3[1] + + K02[3] * E3[2] + + K03[3] * E3[3] + + K04[3] * E3[4] + + K05[3] * E3[5] + + K06[3] * E3[6] + + K07[3] * E3[7] + + K08[3] * E3[8] + + K09[3] * E3[9] + + K10[3] * E3[10] + + K11[3] * E3[11] + + K12[3] * E3[12] + ) + / scale_v[0], + ( + K00[4] * E3[0] + + K01[4] * E3[1] + + K02[4] * E3[2] + + K03[4] * E3[3] + + K04[4] * E3[4] + + K05[4] * E3[5] + + K06[4] * E3[6] + + K07[4] * E3[7] + + K08[4] * E3[8] + + K09[4] * E3[9] + + K10[4] * E3[10] + + K11[4] * E3[11] + + K12[4] * E3[12] + ) + / scale_v[1], + ( + K00[5] * E3[0] + + K01[5] * E3[1] + + K02[5] * E3[2] + + K03[5] * E3[3] + + K04[5] * E3[4] + + K05[5] * E3[5] + + K06[5] * E3[6] + + K07[5] * E3[7] + + K08[5] * E3[8] + + K09[5] * E3[9] + + K10[5] * E3[10] + + K11[5] * E3[11] + + K12[5] * E3[12] + ) + / scale_v[2], + ) + err5 = ( + ( + K00[0] * E5[0] + + K01[0] * E5[1] + + K02[0] * E5[2] + + K03[0] * E5[3] + + K04[0] * E5[4] + + K05[0] * E5[5] + + K06[0] * E5[6] + + K07[0] * E5[7] + + K08[0] * E5[8] + + K09[0] * E5[9] + + K10[0] * E5[10] + + K11[0] * E5[11] + + K12[0] * E5[12] + ) + / scale_r[0], + ( + K00[1] * E5[0] + + K01[1] * E5[1] + + K02[1] * E5[2] + + K03[1] * E5[3] + + K04[1] * E5[4] + + K05[1] * E5[5] + + K06[1] * E5[6] + + K07[1] * E5[7] + + K08[1] * E5[8] + + K09[1] * E5[9] + + K10[1] * E5[10] + + K11[1] * E5[11] + + K12[1] * E5[12] + ) + / scale_r[1], + ( + K00[2] * E5[0] + + K01[2] * E5[1] + + K02[2] * E5[2] + + K03[2] * E5[3] + + K04[2] * E5[4] + + K05[2] * E5[5] + + K06[2] * E5[6] + + K07[2] * E5[7] + + K08[2] * E5[8] + + K09[2] * E5[9] + + K10[2] * E5[10] + + K11[2] * E5[11] + + K12[2] * E5[12] + ) + / scale_r[2], + ( + K00[3] * E5[0] + + K01[3] * E5[1] + + K02[3] * E5[2] + + K03[3] * E5[3] + + K04[3] * E5[4] + + K05[3] * E5[5] + + K06[3] * E5[6] + + K07[3] * E5[7] + + K08[3] * E5[8] + + K09[3] * E5[9] + + K10[3] * E5[10] + + K11[3] * E5[11] + + K12[3] * E5[12] + ) + / scale_v[0], + ( + K00[4] * E5[0] + + K01[4] * E5[1] + + K02[4] * E5[2] + + K03[4] * E5[3] + + K04[4] * E5[4] + + K05[4] * E5[5] + + K06[4] * E5[6] + + K07[4] * E5[7] + + K08[4] * E5[8] + + K09[4] * E5[9] + + K10[4] * E5[10] + + K11[4] * E5[11] + + K12[4] * E5[12] + ) + / scale_v[1], + ( + K00[5] * E5[0] + + K01[5] * E5[1] + + K02[5] * E5[2] + + K03[5] * E5[3] + + K04[5] * E5[4] + + K05[5] * E5[5] + + K06[5] * E5[6] + + K07[5] * E5[7] + + K08[5] * E5[8] + + K09[5] * E5[9] + + K10[5] * E5[10] + + K11[5] * E5[11] + + K12[5] * E5[12] + ) + / scale_v[2], + ) + + err5_norm_2 = ( + err5[0] ** 2 + + err5[1] ** 2 + + err5[2] ** 2 + + err5[3] ** 2 + + err5[4] ** 2 + + err5[5] ** 2 + ) + err3_norm_2 = ( + err3[0] ** 2 + + err3[1] ** 2 + + err3[2] ** 2 + + err3[3] ** 2 + + err3[4] ** 2 + + err3[5] ** 2 + ) + + if err5_norm_2 == 0 and err3_norm_2 == 0: + return 0.0 + denom = err5_norm_2 + 0.01 * err3_norm_2 + + return abs(h) * err5_norm_2 / sqrt(denom * N_RV) diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index 382bdacf6..83d5aa8f1 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -17,6 +17,7 @@ "rk_step_hf", "N_RV", "N_STAGES", + "KSIG", ] N_RV = 6 @@ -36,14 +37,14 @@ B = tuple(float_(number) for number in _B) C = tuple(float_(number) for number in _C[:N_STAGES]) -_KSIG = ( +KSIG = ( "Tuple([" + ",".join(["Tuple([" + ",".join(["f"] * N_RV) + "])"] * (N_STAGES + 1)) + "])" ) -@hjit(f"Tuple([V,V,V,V,{_KSIG:s}])(F({DSIG:s}),f,V,V,V,V,f,f)") +@hjit(f"Tuple([V,V,V,V,{KSIG:s}])(F({DSIG:s}),f,V,V,V,V,f,f)") def rk_step_hf(fun, t, rr, vv, fr, fv, h, argk): """Perform a single Runge-Kutta step. From c30450a9ec82e14eb2a44639d07d6ceb0571a84d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 14:30:42 +0100 Subject: [PATCH 165/346] rm msg --- src/hapsira/core/math/ivp/_api.py | 13 +------------ src/hapsira/core/math/ivp/_rk.py | 11 +++-------- 2 files changed, 4 insertions(+), 20 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 348ead74f..343502d57 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -6,12 +6,6 @@ from ._rk import EPS, DOP853 -MESSAGES = { - 0: "The solver successfully reached the end of the integration interval.", - 1: "A termination event occurred.", -} - - class OdeResult(dict): """Represents the optimization result. @@ -22,8 +16,6 @@ class OdeResult(dict): status : int Termination status of the optimizer. Its value depends on the underlying solver. Refer to `message` for details. - message : str - Description of the cause of the termination. t, y, sol, t_events, y_events, nlu : ? """ @@ -314,8 +306,6 @@ def solve_ivp( * 0: The solver successfully reached the end of `tspan`. * 1: A termination event occurred. - message : string - Human-readable description of the termination reason. success : bool True if the solver reached the interval end or a termination event occurred (``status >= 0``). @@ -344,7 +334,7 @@ def solve_ivp( status = None while status is None: - message = solver.step() + solver.step() if solver.status == "finished": status = 0 @@ -401,6 +391,5 @@ def solve_ivp( y_events=y_events, nlu=solver.nlu, status=status, - message=MESSAGES.get(status, message), success=status >= 0, ) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index c032ccaba..e5c15c1e8 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -218,8 +218,6 @@ class DOP853: Number of LU decompositions. Is always 0 for this solver. """ - TOO_SMALL_STEP = "Required step size is less than spacing between numbers." - A_EXTRA = _A[N_STAGES + 1 :] C_EXTRA = _C[N_STAGES + 1 :] D = _D @@ -305,11 +303,10 @@ def step(self): # Handle corner cases of empty solver or no integration. self.t_old = self.t self.t = self.t_bound - message = None self.status = "finished" else: t = self.t - success, message = self._step_impl() + success = self._step_impl() if not success: self.status = "failed" @@ -318,8 +315,6 @@ def step(self): if self.direction * (self.t - self.t_bound) >= 0: self.status = "finished" - return message - def dense_output(self): """Compute a local interpolant over the last successful step. @@ -379,7 +374,7 @@ def _step_impl(self): while not step_accepted: if h_abs < min_step: - return False, self.TOO_SMALL_STEP + return False h = h_abs * self.direction t_new = t + h @@ -438,4 +433,4 @@ def _step_impl(self): self.h_abs = h_abs self.f = f_new - return True, None + return True From 09914f6a3ad871cdfec8b49fae6ae63b06294324 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 14:40:04 +0100 Subject: [PATCH 166/346] cleanup --- src/hapsira/core/math/ivp/_rk.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index e5c15c1e8..eb8cb3b3e 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -304,16 +304,20 @@ def step(self): self.t_old = self.t self.t = self.t_bound self.status = "finished" - else: - t = self.t - success = self._step_impl() + return - if not success: - self.status = "failed" - else: - self.t_old = t - if self.direction * (self.t - self.t_bound) >= 0: - self.status = "finished" + t = self.t + success = self._step_impl() + + if not success: + self.status = "failed" + return + + self.t_old = t + if self.direction * (self.t - self.t_bound) < 0: + return + + self.status = "finished" def dense_output(self): """Compute a local interpolant over the last successful step. From 350bf67badcf6159f8c16cbc99abe8f3ea37b001 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 14:53:25 +0100 Subject: [PATCH 167/346] cleanup --- src/hapsira/core/math/ivp/_api.py | 1 - src/hapsira/core/math/ivp/_rk.py | 120 ++++++------------------------ 2 files changed, 21 insertions(+), 100 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 343502d57..d6a9e0ff4 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -389,7 +389,6 @@ def solve_ivp( sol=sol, t_events=t_events, y_events=y_events, - nlu=solver.nlu, status=status, success=status >= 0, ) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index eb8cb3b3e..416f4d6e9 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -27,6 +27,7 @@ INTERPOLATOR_POWER = 7 N_STAGES_EXTENDED = 16 ERROR_ESTIMATOR_ORDER = 7 +ERROR_EXPONENT = -1 / (ERROR_ESTIMATOR_ORDER + 1) @hjit("f(V,V)") @@ -138,84 +139,9 @@ def __call__(self, t): return y.T -def check_arguments(y0): - """Helper function for checking arguments common to all solvers.""" - - y0 = np.asarray(y0) - - assert not np.issubdtype(y0.dtype, np.complexfloating) - - dtype = float - y0 = y0.astype(dtype, copy=False) - - assert not y0.ndim != 1 - assert np.isfinite(y0).all() - - return y0 - - class DOP853: - """Explicit Runge-Kutta method of order 8. - - This is a Python implementation of "DOP853" algorithm originally written - in Fortran [1]_, [2]_. Note that this is not a literate translation, but - the algorithmic core and coefficients are the same. - - Parameters - ---------- - fun : callable - Right-hand side of the system. The calling signature is ``fun(t, y)``. - Here, ``t`` is a scalar, and there are two options for the ndarray ``y``: - It can either have shape (n,); then ``fun`` must return array_like with - shape (n,). Alternatively it can have shape (n, k); then ``fun`` - must return an array_like with shape (n, k), i.e. each column - corresponds to a single column in ``y``. The choice between the two - options is determined by `vectorized` argument (see below). - t0 : float - Initial time. - y0 : array_like, shape (n,) - Initial state. - t_bound : float - Boundary time - the integration won't continue beyond it. It also - determines the direction of the integration. - max_step : float, optional - Maximum allowed step size. Default is np.inf, i.e. the step size is not - bounded and determined solely by the solver. - rtol, atol : float and array_like, optional - Relative and absolute tolerances. The solver keeps the local error - estimates less than ``atol + rtol * abs(y)``. Here `rtol` controls a - relative accuracy (number of correct digits), while `atol` controls - absolute accuracy (number of correct decimal places). To achieve the - desired `rtol`, set `atol` to be smaller than the smallest value that - can be expected from ``rtol * abs(y)`` so that `rtol` dominates the - allowable error. If `atol` is larger than ``rtol * abs(y)`` the - number of correct digits is not guaranteed. Conversely, to achieve the - desired `atol` set `rtol` such that ``rtol * abs(y)`` is always smaller - than `atol`. If components of y have different scales, it might be - beneficial to set different `atol` values for different components by - passing array_like with shape (n,) for `atol`. Default values are - 1e-3 for `rtol` and 1e-6 for `atol`. - - Attributes - ---------- - n : int - Number of equations. - status : string - Current status of the solver: 'running', 'finished' or 'failed'. - t_bound : float - Boundary time. - direction : float - Integration direction: +1 or -1. - t : float - Current time. - y : ndarray - Current state. - t_old : float - Previous time. None if no steps were made yet. - step_size : float - Size of the last successful step. None if no steps were made yet. - nlu : int - Number of LU decompositions. Is always 0 for this solver. + """ + Explicit Runge-Kutta method of order 8. """ A_EXTRA = _A[N_STAGES + 1 :] @@ -234,31 +160,31 @@ def __init__( atol: float = 1e-6, ): assert y0.shape == (N_RV,) + assert np.isfinite(y0).all() + assert max_step > 0 + assert atol >= 0 + + if rtol < 100 * EPS: + rtol = 100 * EPS - self.t_old = None self.t = t0 - self.y = check_arguments(y0) + self.y = y0 self.t_bound = t_bound - + self.max_step = max_step self.fun = fun self.argk = argk + self.rtol = rtol + self.atol = atol self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 - self.status = "running" - - self.nfev = 0 - self.njev = 0 - self.nlu = 0 + self.K_extended = np.empty((N_STAGES_EXTENDED, N_RV), dtype=self.y.dtype) + self.K = self.K_extended[: N_STAGES + 1, :] self.y_old = None + self.t_old = None + self.h_previous = None - assert max_step > 0 - self.max_step = max_step - - if rtol < 100 * EPS: - rtol = 100 * EPS - assert atol >= 0 - self.rtol, self.atol = rtol, atol + self.status = "running" rr, vv = self.fun( self.t, @@ -267,6 +193,7 @@ def __init__( self.argk, ) # TODO call into hf self.f = np.array([*rr, *vv]) + self.h_abs = _select_initial_step_hf( self.fun, self.t, @@ -280,11 +207,6 @@ def __init__( self.rtol, self.atol, ) # TODO call into hf - self.error_exponent = -1 / (ERROR_ESTIMATOR_ORDER + 1) - self.h_previous = None - - self.K_extended = np.empty((N_STAGES_EXTENDED, N_RV), dtype=self.y.dtype) - self.K = self.K_extended[: N_STAGES + 1] def step(self): """Perform one integration step. @@ -416,7 +338,7 @@ def _step_impl(self): if error_norm == 0: factor = MAX_FACTOR else: - factor = min(MAX_FACTOR, SAFETY * error_norm**self.error_exponent) + factor = min(MAX_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) if step_rejected: factor = min(1, factor) @@ -425,7 +347,7 @@ def _step_impl(self): step_accepted = True else: - h_abs *= max(MIN_FACTOR, SAFETY * error_norm**self.error_exponent) + h_abs *= max(MIN_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) step_rejected = True self.h_previous = h From 0f38b9dea46116a392e7e73a73d042701cfd56a0 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 20:04:31 +0100 Subject: [PATCH 168/346] more linalg --- src/hapsira/core/math/linalg.py | 40 +++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 7c6ea5092..3e9454ffe 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -1,8 +1,11 @@ from math import inf, sqrt +from numpy import finfo from ..jit import hjit, vjit +from ...settings import settings __all__ = [ + "add_Vs_hf", "add_VV_hf", "cross_VV_hf", "div_Vs_hf", @@ -10,16 +13,37 @@ "matmul_MM_hf", "matmul_VM_hf", "matmul_VV_hf", + "max_VV_hf", "mul_Vs_hf", "mul_VV_hf", + "nextafter_hf", "norm_hf", "norm_vf", "sign_hf", "sub_VV_hf", "transpose_M_hf", + "EPS", ] +if settings["PRECISION"].value == "f8": + from numpy import float64 as float_ +elif settings["PRECISION"].value == "f4": + from numpy import float32 as float_ +elif settings["PRECISION"].value == "f2": + from numpy import float16 as float_ +else: + raise ValueError("unsupported precision") + + +EPS = finfo(float_).eps + + +@hjit("V(V,f)") +def add_Vs_hf(a, b): + return a[0] + b, a[1] + b, a[2] + b + + @hjit("V(V,V)") def add_VV_hf(a, b): return a[0] + b[0], a[1] + b[1], a[2] + b[2] @@ -89,6 +113,15 @@ def matmul_VV_hf(a, b): return a[0] * b[0] + a[1] * b[1] + a[2] * b[2] +@hjit("V(V,V)") +def max_VV_hf(x, y): + return ( + x[0] if x[0] > y[0] else y[0], + x[1] if x[1] > y[1] else y[1], + x[2] if x[2] > y[2] else y[2], + ) + + @hjit("V(V,f)") def mul_Vs_hf(v, s): return v[0] * s, v[1] * s, v[2] * s @@ -99,6 +132,13 @@ def mul_VV_hf(a, b): return a[0] * b[0], a[1] * b[1], a[2] * b[2] +@hjit("f(f,f)") +def nextafter_hf(x, direction): + if x < direction: + return x + EPS + return x - EPS + + @hjit("f(V)") def norm_hf(a): return sqrt(matmul_VV_hf(a, a)) From 7babf7f0df57ae5d1f1d4d6e72ac3e646cad409d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 20:09:28 +0100 Subject: [PATCH 169/346] move eps --- src/hapsira/core/math/ivp/_api.py | 4 ++-- src/hapsira/core/math/ivp/_rk.py | 5 +---- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index d6a9e0ff4..77ddeeaa2 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -2,8 +2,8 @@ from ._brentq import brentq from ._common import OdeSolution - -from ._rk import EPS, DOP853 +from ._rk import DOP853 +from ...math.linalg import EPS class OdeResult(dict): diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 416f4d6e9..4cd7b0c0f 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -8,16 +8,13 @@ from ._rkerror import estimate_error_norm_hf from ...jit import array_to_V_hf, hjit, DSIG -from ...math.linalg import add_VV_hf, div_VV_hf, mul_Vs_hf, sub_VV_hf +from ...math.linalg import add_VV_hf, div_VV_hf, mul_Vs_hf, sub_VV_hf, EPS __all__ = [ - "EPS", "DOP853", ] -EPS = np.finfo(float).eps - # Multiply steps computed from asymptotic behaviour of errors by this. SAFETY = 0.9 From ce449796d453d8918baed4d1a53e9ff8b5e84247 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 20:12:53 +0100 Subject: [PATCH 170/346] rm numpy dep --- src/hapsira/core/math/ivp/_rk.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 4cd7b0c0f..d0ee7656e 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -1,4 +1,4 @@ -from math import sqrt +from math import inf, sqrt from typing import Callable import numpy as np @@ -8,7 +8,7 @@ from ._rkerror import estimate_error_norm_hf from ...jit import array_to_V_hf, hjit, DSIG -from ...math.linalg import add_VV_hf, div_VV_hf, mul_Vs_hf, sub_VV_hf, EPS +from ...math.linalg import add_VV_hf, div_VV_hf, mul_Vs_hf, nextafter_hf, sub_VV_hf, EPS __all__ = [ "DOP853", @@ -283,7 +283,7 @@ def _step_impl(self): rtol = self.rtol atol = self.atol - min_step = 10 * np.abs(np.nextafter(t, self.direction * np.inf) - t) + min_step = 10 * abs(nextafter_hf(t, self.direction * inf) - t) if self.h_abs > max_step: h_abs = max_step From 4d753a8ec5803017b6115f6e685c54199f6c48ed Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 20:23:40 +0100 Subject: [PATCH 171/346] less numpy --- src/hapsira/core/math/ivp/_rk.py | 40 +++++++++++++++++++++++++++----- src/hapsira/core/math/linalg.py | 6 +++++ 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index d0ee7656e..9248d74cd 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -8,7 +8,17 @@ from ._rkerror import estimate_error_norm_hf from ...jit import array_to_V_hf, hjit, DSIG -from ...math.linalg import add_VV_hf, div_VV_hf, mul_Vs_hf, nextafter_hf, sub_VV_hf, EPS +from ...math.linalg import ( + abs_V_hf, + add_Vs_hf, + add_VV_hf, + div_VV_hf, + max_VV_hf, + mul_Vs_hf, + nextafter_hf, + sub_VV_hf, + EPS, +) __all__ = [ "DOP853", @@ -306,7 +316,7 @@ def _step_impl(self): t_new = self.t_bound h = t_new - t - h_abs = np.abs(h) + h_abs = abs(h) rr_new, vv_new, fr_new, fv_new, K_new = rk_step_hf( self.fun, @@ -322,13 +332,31 @@ def _step_impl(self): f_new = np.array([*fr_new, *fv_new]) self.K[: N_STAGES + 1, :N_RV] = np.array([K_new]) - scale = atol + np.maximum(np.abs(y), np.abs(y_new)) * rtol - assert scale.shape == (N_RV,) + scale_r = add_Vs_hf( + mul_Vs_hf( + max_VV_hf( + abs_V_hf(array_to_V_hf(y[:3])), + abs_V_hf(array_to_V_hf(y_new[:3])), + ), + rtol, + ), + atol, + ) + scale_v = add_Vs_hf( + mul_Vs_hf( + max_VV_hf( + abs_V_hf(array_to_V_hf(y[3:])), + abs_V_hf(array_to_V_hf(y_new[3:])), + ), + rtol, + ), + atol, + ) error_norm = estimate_error_norm_hf( K_new, h, - array_to_V_hf(scale[:3]), - array_to_V_hf(scale[3:]), + scale_r, + scale_v, ) # TODO call into hf if error_norm < 1: diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 3e9454ffe..4c2e78087 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -5,6 +5,7 @@ from ...settings import settings __all__ = [ + "abs_V_hf", "add_Vs_hf", "add_VV_hf", "cross_VV_hf", @@ -39,6 +40,11 @@ EPS = finfo(float_).eps +@hjit("V(V)") +def abs_V_hf(x): + return abs(x[0]), abs(x[1]), abs(x[2]) + + @hjit("V(V,f)") def add_Vs_hf(a, b): return a[0] + b, a[1] + b, a[2] + b From 19d593f5b43db105b8f0aaa321312412e4b17fa4 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 20:27:26 +0100 Subject: [PATCH 172/346] use vectors --- src/hapsira/core/math/ivp/_rk.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 9248d74cd..4d3ef1fb4 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -336,7 +336,7 @@ def _step_impl(self): mul_Vs_hf( max_VV_hf( abs_V_hf(array_to_V_hf(y[:3])), - abs_V_hf(array_to_V_hf(y_new[:3])), + abs_V_hf(rr_new), ), rtol, ), @@ -346,7 +346,7 @@ def _step_impl(self): mul_Vs_hf( max_VV_hf( abs_V_hf(array_to_V_hf(y[3:])), - abs_V_hf(array_to_V_hf(y_new[3:])), + abs_V_hf(vv_new), ), rtol, ), From 80fd3ceb5e496b1e2a3ddbbfdb9d3363b4577048 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 20:44:13 +0100 Subject: [PATCH 173/346] mv step impl out of class --- src/hapsira/core/math/ivp/_rk.py | 174 ++++++++++++++++--------------- 1 file changed, 89 insertions(+), 85 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 4d3ef1fb4..5176f1af8 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -236,7 +236,29 @@ def step(self): return t = self.t - success = self._step_impl() + success, *rets = _step_impl( + self.fun, + self.argk, + self.t, + self.y, + self.f, + self.max_step, + self.rtol, + self.atol, + self.direction, + self.h_abs, + self.t_bound, + tuple(tuple(line) for line in self.K[: N_STAGES + 1, :N_RV]), + ) + + if success: + self.h_previous = rets[0] + self.y_old = rets[1] + self.t = rets[2] + self.y = rets[3] + self.h_abs = rets[4] + self.f = rets[5] + self.K[: N_STAGES + 1, :N_RV] = np.array(rets[6]) if not success: self.status = "failed" @@ -285,103 +307,85 @@ def dense_output(self): return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) - def _step_impl(self): - t = self.t - y = self.y - max_step = self.max_step - rtol = self.rtol - atol = self.atol +def _step_impl(fun, argk, t, y, f, max_step, rtol, atol, direction, h_abs, t_bound, K): + min_step = 10 * abs(nextafter_hf(t, direction * inf) - t) - min_step = 10 * abs(nextafter_hf(t, self.direction * inf) - t) - - if self.h_abs > max_step: - h_abs = max_step - elif self.h_abs < min_step: - h_abs = min_step - else: - h_abs = self.h_abs + if h_abs > max_step: + h_abs = max_step + if h_abs < min_step: + h_abs = min_step - step_accepted = False - step_rejected = False + step_accepted = False + step_rejected = False - while not step_accepted: - if h_abs < min_step: - return False + while not step_accepted: + if h_abs < min_step: + return False - h = h_abs * self.direction - t_new = t + h + h = h_abs * direction + t_new = t + h - if self.direction * (t_new - self.t_bound) > 0: - t_new = self.t_bound + if direction * (t_new - t_bound) > 0: + t_new = t_bound - h = t_new - t - h_abs = abs(h) + h = t_new - t + h_abs = abs(h) - rr_new, vv_new, fr_new, fv_new, K_new = rk_step_hf( - self.fun, - t, - array_to_V_hf(y[:3]), - array_to_V_hf(y[3:]), - array_to_V_hf(self.f[:3]), - array_to_V_hf(self.f[3:]), - h, - self.argk, - ) # TODO call into hf - y_new = np.array([*rr_new, *vv_new]) - f_new = np.array([*fr_new, *fv_new]) - self.K[: N_STAGES + 1, :N_RV] = np.array([K_new]) - - scale_r = add_Vs_hf( - mul_Vs_hf( - max_VV_hf( - abs_V_hf(array_to_V_hf(y[:3])), - abs_V_hf(rr_new), - ), - rtol, + rr_new, vv_new, fr_new, fv_new, K_new = rk_step_hf( + fun, + t, + array_to_V_hf(y[:3]), + array_to_V_hf(y[3:]), + array_to_V_hf(f[:3]), + array_to_V_hf(f[3:]), + h, + argk, + ) + y_new = np.array([*rr_new, *vv_new]) + f_new = np.array([*fr_new, *fv_new]) + + scale_r = add_Vs_hf( + mul_Vs_hf( + max_VV_hf( + abs_V_hf(array_to_V_hf(y[:3])), + abs_V_hf(rr_new), ), - atol, - ) - scale_v = add_Vs_hf( - mul_Vs_hf( - max_VV_hf( - abs_V_hf(array_to_V_hf(y[3:])), - abs_V_hf(vv_new), - ), - rtol, + rtol, + ), + atol, + ) + scale_v = add_Vs_hf( + mul_Vs_hf( + max_VV_hf( + abs_V_hf(array_to_V_hf(y[3:])), + abs_V_hf(vv_new), ), - atol, - ) - error_norm = estimate_error_norm_hf( - K_new, - h, - scale_r, - scale_v, - ) # TODO call into hf - - if error_norm < 1: - if error_norm == 0: - factor = MAX_FACTOR - else: - factor = min(MAX_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) - - if step_rejected: - factor = min(1, factor) - - h_abs *= factor + rtol, + ), + atol, + ) + error_norm = estimate_error_norm_hf( + K_new, + h, + scale_r, + scale_v, + ) - step_accepted = True + if error_norm < 1: + if error_norm == 0: + factor = MAX_FACTOR else: - h_abs *= max(MIN_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) - step_rejected = True + factor = min(MAX_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) - self.h_previous = h - self.y_old = y + if step_rejected: + factor = min(1, factor) - self.t = t_new - self.y = y_new + h_abs *= factor - self.h_abs = h_abs - self.f = f_new + step_accepted = True + else: + h_abs *= max(MIN_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) + step_rejected = True - return True + return True, h, y, t_new, y_new, h_abs, f_new, K_new From b30315c9657cfd7cb80cb585c9bd064ec77f397a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 21:20:43 +0100 Subject: [PATCH 174/346] f to fr fv --- src/hapsira/core/math/ivp/_rk.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 5176f1af8..26bc2c936 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -241,7 +241,8 @@ def step(self): self.argk, self.t, self.y, - self.f, + array_to_V_hf(self.f[:3]), + array_to_V_hf(self.f[3:]), self.max_step, self.rtol, self.atol, @@ -257,8 +258,8 @@ def step(self): self.t = rets[2] self.y = rets[3] self.h_abs = rets[4] - self.f = rets[5] - self.K[: N_STAGES + 1, :N_RV] = np.array(rets[6]) + self.f = np.array([*rets[5], *rets[6]]) + self.K[: N_STAGES + 1, :N_RV] = np.array(rets[7]) if not success: self.status = "failed" @@ -308,7 +309,9 @@ def dense_output(self): return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) -def _step_impl(fun, argk, t, y, f, max_step, rtol, atol, direction, h_abs, t_bound, K): +def _step_impl( + fun, argk, t, y, fr, fv, max_step, rtol, atol, direction, h_abs, t_bound, K +): min_step = 10 * abs(nextafter_hf(t, direction * inf) - t) if h_abs > max_step: @@ -337,13 +340,12 @@ def _step_impl(fun, argk, t, y, f, max_step, rtol, atol, direction, h_abs, t_bou t, array_to_V_hf(y[:3]), array_to_V_hf(y[3:]), - array_to_V_hf(f[:3]), - array_to_V_hf(f[3:]), + fr, + fv, h, argk, ) y_new = np.array([*rr_new, *vv_new]) - f_new = np.array([*fr_new, *fv_new]) scale_r = add_Vs_hf( mul_Vs_hf( @@ -388,4 +390,4 @@ def _step_impl(fun, argk, t, y, f, max_step, rtol, atol, direction, h_abs, t_bou h_abs *= max(MIN_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) step_rejected = True - return True, h, y, t_new, y_new, h_abs, f_new, K_new + return True, h, y, t_new, y_new, h_abs, fr_new, fv_new, K_new From 54ddf59efd2f47a5a2b73e4bb225552b986ff76e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 21:25:58 +0100 Subject: [PATCH 175/346] y to rr vv --- src/hapsira/core/math/ivp/_rk.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 26bc2c936..441d5ed7c 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -240,7 +240,8 @@ def step(self): self.fun, self.argk, self.t, - self.y, + array_to_V_hf(self.y[:3]), + array_to_V_hf(self.y[3:]), array_to_V_hf(self.f[:3]), array_to_V_hf(self.f[3:]), self.max_step, @@ -254,12 +255,12 @@ def step(self): if success: self.h_previous = rets[0] - self.y_old = rets[1] - self.t = rets[2] - self.y = rets[3] - self.h_abs = rets[4] - self.f = np.array([*rets[5], *rets[6]]) - self.K[: N_STAGES + 1, :N_RV] = np.array(rets[7]) + self.y_old = np.array([*rets[1], *rets[2]]) + self.t = rets[3] + self.y = np.array([*rets[4], *rets[5]]) + self.h_abs = rets[6] + self.f = np.array([*rets[7], *rets[8]]) + self.K[: N_STAGES + 1, :N_RV] = np.array(rets[9]) if not success: self.status = "failed" @@ -310,7 +311,7 @@ def dense_output(self): def _step_impl( - fun, argk, t, y, fr, fv, max_step, rtol, atol, direction, h_abs, t_bound, K + fun, argk, t, rr, vv, fr, fv, max_step, rtol, atol, direction, h_abs, t_bound, K ): min_step = 10 * abs(nextafter_hf(t, direction * inf) - t) @@ -338,19 +339,18 @@ def _step_impl( rr_new, vv_new, fr_new, fv_new, K_new = rk_step_hf( fun, t, - array_to_V_hf(y[:3]), - array_to_V_hf(y[3:]), + rr, + vv, fr, fv, h, argk, ) - y_new = np.array([*rr_new, *vv_new]) scale_r = add_Vs_hf( mul_Vs_hf( max_VV_hf( - abs_V_hf(array_to_V_hf(y[:3])), + abs_V_hf(rr), abs_V_hf(rr_new), ), rtol, @@ -360,7 +360,7 @@ def _step_impl( scale_v = add_Vs_hf( mul_Vs_hf( max_VV_hf( - abs_V_hf(array_to_V_hf(y[3:])), + abs_V_hf(vv), abs_V_hf(vv_new), ), rtol, @@ -390,4 +390,4 @@ def _step_impl( h_abs *= max(MIN_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) step_rejected = True - return True, h, y, t_new, y_new, h_abs, fr_new, fv_new, K_new + return True, h, rr, vv, t_new, rr_new, vv_new, h_abs, fr_new, fv_new, K_new From 04bd4bb3de48114874cf31e179ebc41a1076c607 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 21:29:36 +0100 Subject: [PATCH 176/346] side-effect free step impl --- src/hapsira/core/math/ivp/_rk.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 441d5ed7c..45b67f88c 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -255,12 +255,13 @@ def step(self): if success: self.h_previous = rets[0] - self.y_old = np.array([*rets[1], *rets[2]]) - self.t = rets[3] - self.y = np.array([*rets[4], *rets[5]]) - self.h_abs = rets[6] - self.f = np.array([*rets[7], *rets[8]]) - self.K[: N_STAGES + 1, :N_RV] = np.array(rets[9]) + # self.y_old = np.array([*rets[1], *rets[2]]) + self.y_old = self.y + self.t = rets[1] + self.y = np.array([*rets[2], *rets[3]]) + self.h_abs = rets[4] + self.f = np.array([*rets[5], *rets[6]]) + self.K[: N_STAGES + 1, :N_RV] = np.array(rets[7]) if not success: self.status = "failed" @@ -390,4 +391,4 @@ def _step_impl( h_abs *= max(MIN_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) step_rejected = True - return True, h, rr, vv, t_new, rr_new, vv_new, h_abs, fr_new, fv_new, K_new + return True, h, t_new, rr_new, vv_new, h_abs, fr_new, fv_new, K_new From e317130a3c7275dec68ca0aa840da7897598fed6 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 21:36:12 +0100 Subject: [PATCH 177/346] jit step impl --- src/hapsira/core/math/ivp/_rk.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 45b67f88c..875d4b789 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -4,7 +4,7 @@ import numpy as np from ._dop853_coefficients import A as _A, C as _C, D as _D -from ._rkstep import rk_step_hf, N_RV, N_STAGES +from ._rkstep import rk_step_hf, N_RV, N_STAGES, KSIG from ._rkerror import estimate_error_norm_hf from ...jit import array_to_V_hf, hjit, DSIG @@ -236,7 +236,7 @@ def step(self): return t = self.t - success, *rets = _step_impl( + success, *rets = _step_impl_hf( self.fun, self.argk, self.t, @@ -311,7 +311,11 @@ def dense_output(self): return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) -def _step_impl( +@hjit( + f"Tuple([b1,f,f,V,V,f,V,V,{KSIG:s}])" + f"(F({DSIG:s}),f,f,V,V,V,V,f,f,f,f,f,f,{KSIG:s})" +) +def _step_impl_hf( fun, argk, t, rr, vv, fr, fv, max_step, rtol, atol, direction, h_abs, t_bound, K ): min_step = 10 * abs(nextafter_hf(t, direction * inf) - t) @@ -326,7 +330,17 @@ def _step_impl( while not step_accepted: if h_abs < min_step: - return False + return ( + False, + 0.0, + 0.0, + (0.0, 0.0, 0.0), + (0.0, 0.0, 0.0), + 0.0, + (0.0, 0.0, 0.0), + (0.0, 0.0, 0.0), + K, + ) h = h_abs * direction t_new = t + h From 55e057cb284299c5a6da9bf1bbe4e837ddd71dc3 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 21:43:09 +0100 Subject: [PATCH 178/346] mv function --- src/hapsira/core/math/ivp/_rk.py | 194 +++++++++++++++---------------- 1 file changed, 97 insertions(+), 97 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 875d4b789..6a882dd65 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -92,6 +92,103 @@ def _select_initial_step_hf( return min(100 * h0, h1) +@hjit( + f"Tuple([b1,f,f,V,V,f,V,V,{KSIG:s}])" + f"(F({DSIG:s}),f,f,V,V,V,V,f,f,f,f,f,f,{KSIG:s})" +) +def _step_impl_hf( + fun, argk, t, rr, vv, fr, fv, max_step, rtol, atol, direction, h_abs, t_bound, K +): + min_step = 10 * abs(nextafter_hf(t, direction * inf) - t) + + if h_abs > max_step: + h_abs = max_step + if h_abs < min_step: + h_abs = min_step + + step_accepted = False + step_rejected = False + + while not step_accepted: + if h_abs < min_step: + return ( + False, + 0.0, + 0.0, + (0.0, 0.0, 0.0), + (0.0, 0.0, 0.0), + 0.0, + (0.0, 0.0, 0.0), + (0.0, 0.0, 0.0), + K, + ) + + h = h_abs * direction + t_new = t + h + + if direction * (t_new - t_bound) > 0: + t_new = t_bound + + h = t_new - t + h_abs = abs(h) + + rr_new, vv_new, fr_new, fv_new, K_new = rk_step_hf( + fun, + t, + rr, + vv, + fr, + fv, + h, + argk, + ) + + scale_r = add_Vs_hf( + mul_Vs_hf( + max_VV_hf( + abs_V_hf(rr), + abs_V_hf(rr_new), + ), + rtol, + ), + atol, + ) + scale_v = add_Vs_hf( + mul_Vs_hf( + max_VV_hf( + abs_V_hf(vv), + abs_V_hf(vv_new), + ), + rtol, + ), + atol, + ) + error_norm = estimate_error_norm_hf( + K_new, + h, + scale_r, + scale_v, + ) + + if error_norm < 1: + if error_norm == 0: + factor = MAX_FACTOR + else: + factor = min(MAX_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) + + if step_rejected: + factor = min(1, factor) + + h_abs *= factor + + step_accepted = True + else: + h_abs *= max(MIN_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) + step_rejected = True + + return True, h, t_new, rr_new, vv_new, h_abs, fr_new, fv_new, K_new + + class Dop853DenseOutput: """local interpolant over step made by an ODE solver. @@ -309,100 +406,3 @@ def dense_output(self): F[3:] = h * np.dot(self.D, K) return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) - - -@hjit( - f"Tuple([b1,f,f,V,V,f,V,V,{KSIG:s}])" - f"(F({DSIG:s}),f,f,V,V,V,V,f,f,f,f,f,f,{KSIG:s})" -) -def _step_impl_hf( - fun, argk, t, rr, vv, fr, fv, max_step, rtol, atol, direction, h_abs, t_bound, K -): - min_step = 10 * abs(nextafter_hf(t, direction * inf) - t) - - if h_abs > max_step: - h_abs = max_step - if h_abs < min_step: - h_abs = min_step - - step_accepted = False - step_rejected = False - - while not step_accepted: - if h_abs < min_step: - return ( - False, - 0.0, - 0.0, - (0.0, 0.0, 0.0), - (0.0, 0.0, 0.0), - 0.0, - (0.0, 0.0, 0.0), - (0.0, 0.0, 0.0), - K, - ) - - h = h_abs * direction - t_new = t + h - - if direction * (t_new - t_bound) > 0: - t_new = t_bound - - h = t_new - t - h_abs = abs(h) - - rr_new, vv_new, fr_new, fv_new, K_new = rk_step_hf( - fun, - t, - rr, - vv, - fr, - fv, - h, - argk, - ) - - scale_r = add_Vs_hf( - mul_Vs_hf( - max_VV_hf( - abs_V_hf(rr), - abs_V_hf(rr_new), - ), - rtol, - ), - atol, - ) - scale_v = add_Vs_hf( - mul_Vs_hf( - max_VV_hf( - abs_V_hf(vv), - abs_V_hf(vv_new), - ), - rtol, - ), - atol, - ) - error_norm = estimate_error_norm_hf( - K_new, - h, - scale_r, - scale_v, - ) - - if error_norm < 1: - if error_norm == 0: - factor = MAX_FACTOR - else: - factor = min(MAX_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) - - if step_rejected: - factor = min(1, factor) - - h_abs *= factor - - step_accepted = True - else: - h_abs *= max(MIN_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) - step_rejected = True - - return True, h, t_new, rr_new, vv_new, h_abs, fr_new, fv_new, K_new From bc319cbb6fe520465ae6845346211234f85863f0 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 21:59:26 +0100 Subject: [PATCH 179/346] cleanup --- src/hapsira/core/math/ivp/_api.py | 20 +++++++++--------- src/hapsira/core/propagation/cowell.py | 28 ++++++++++++++------------ 2 files changed, 24 insertions(+), 24 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 77ddeeaa2..86a9ef497 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -1,3 +1,5 @@ +from typing import Callable, List, Optional + import numpy as np from ._brentq import brentq @@ -160,16 +162,14 @@ def find_active_events(g, g_new, direction): def solve_ivp( - fun, - t_span, - y0, + fun: Callable, + t0: float, + tf: float, + y0: np.ndarray, # (6,) argk: float, - method=DOP853, - # t_eval=None, - # dense_output=False, - events=None, + events: Optional[List[Callable]] = None, **options, -): +) -> OdeResult: """Solve an initial value problem for a system of ODEs. Parameters @@ -312,9 +312,7 @@ def solve_ivp( """ - t0, tf = map(float, t_span) - - solver = method(fun, t0, y0, tf, argk, **options) + solver = DOP853(fun, t0, y0, tf, argk, **options) ts = [t0] ys = [y0] diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 6a03adb46..0ed73c645 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -9,16 +9,20 @@ ] -def cowell(k, r, v, tofs, rtol=1e-11, events=None, f=func_twobody_hf): +def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=None, f=func_twobody_hf): """ Scalar cowell - f : float + k : float r : ndarray (3,) v : ndarray (3,) - tofs : ??? - rtol : float ... or also ndarray? + tofs : array of relative times [seconds] + rtol : float + atol : float + events : Optional[List[Event]] + f : Callable """ + assert hasattr(f, "djit") # DEBUG check for compiler flag assert isinstance(rtol, float) @@ -29,12 +33,12 @@ def cowell(k, r, v, tofs, rtol=1e-11, events=None, f=func_twobody_hf): result = solve_ivp( f, - (0, max(tofs)), + 0.0, + float(max(tofs)), u0, argk=k, rtol=rtol, - atol=1e-12, - # dense_output=True, + atol=atol, events=events, ) if not result.success: @@ -46,18 +50,16 @@ def cowell(k, r, v, tofs, rtol=1e-11, events=None, f=func_twobody_hf): # If there are no terminal events, then the last time of integration is the # greatest one from the original array of propagation times - if not terminal_events: - last_t = max(tofs) - else: + if terminal_events: # Filter the event which triggered first last_t = min(event._last_t for event in terminal_events) # FIXME: Here last_t has units, but tofs don't - tofs = [tof for tof in tofs if tof < last_t] + [last_t] + tofs = [tof for tof in tofs if tof < last_t] + tofs.append(last_t) rrs = [] vvs = [] - for i in range(len(tofs)): - t = tofs[i] + for t in tofs: y = result.sol(t) rrs.append(y[:3]) vvs.append(y[3:]) From 81de503a4ae24ff4de2d885c5cb6e073175ef33e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 22:04:56 +0100 Subject: [PATCH 180/346] cleanup --- src/hapsira/core/math/ivp/_api.py | 40 ++------------------------ src/hapsira/core/propagation/cowell.py | 6 ++-- 2 files changed, 6 insertions(+), 40 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 86a9ef497..2d3409d48 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -1,4 +1,4 @@ -from typing import Callable, List, Optional +from typing import Callable, List, Optional, Tuple import numpy as np @@ -8,32 +8,6 @@ from ...math.linalg import EPS -class OdeResult(dict): - """Represents the optimization result. - - Attributes - ---------- - success : bool - Whether or not the optimizer exited successfully. - status : int - Termination status of the optimizer. Its value depends on the - underlying solver. Refer to `message` for details. - t, y, sol, t_events, y_events, nlu : ? - """ - - def __getattr__(self, name): - try: - return self[name] - except KeyError as e: - raise AttributeError(name) from e - - __setattr__ = dict.__setitem__ - __delattr__ = dict.__delitem__ - - def __dir__(self): - return list(self.keys()) - - def solve_event_equation(event, sol, t_old, t): """Solve an equation corresponding to an ODE event. @@ -169,7 +143,7 @@ def solve_ivp( argk: float, events: Optional[List[Callable]] = None, **options, -) -> OdeResult: +) -> Tuple[OdeSolution, bool]: """Solve an initial value problem for a system of ODEs. Parameters @@ -381,12 +355,4 @@ def solve_ivp( sol = OdeSolution(ts, interpolants) - return OdeResult( - t=ts, - y=ys, - sol=sol, - t_events=t_events, - y_events=y_events, - status=status, - success=status >= 0, - ) + return sol, status >= 0 diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 0ed73c645..be5dd45e9 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -31,7 +31,7 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=None, f=func_twobody_hf u0 = np.array([x, y, z, vx, vy, vz]) - result = solve_ivp( + sol, success = solve_ivp( f, 0.0, float(max(tofs)), @@ -41,7 +41,7 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=None, f=func_twobody_hf atol=atol, events=events, ) - if not result.success: + if not success: raise RuntimeError("Integration failed") if events is not None: @@ -60,7 +60,7 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=None, f=func_twobody_hf rrs = [] vvs = [] for t in tofs: - y = result.sol(t) + y = sol(t) rrs.append(y[:3]) vvs.append(y[3:]) From 73215b6d21bbeea0028263e1df53af4ea93467a8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 22:13:01 +0100 Subject: [PATCH 181/346] cleanup --- src/hapsira/core/math/ivp/_api.py | 13 +------------ 1 file changed, 1 insertion(+), 12 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 2d3409d48..5a9d0a2eb 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -289,7 +289,6 @@ def solve_ivp( solver = DOP853(fun, t0, y0, tf, argk, **options) ts = [t0] - ys = [y0] interpolants = [] @@ -344,15 +343,5 @@ def solve_ivp( g = g_new ts.append(t) - ys.append(y) - if t_events is not None: - t_events = [np.asarray(te) for te in t_events] - y_events = [np.asarray(ye) for ye in y_events] - - ts = np.array(ts) - ys = np.vstack(ys).T - - sol = OdeSolution(ts, interpolants) - - return sol, status >= 0 + return OdeSolution(np.array(ts), interpolants), status >= 0 From 256c4cc403fa370e308ed90403811c7f6e13aba4 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 24 Jan 2024 22:19:38 +0100 Subject: [PATCH 182/346] cleanup --- src/hapsira/core/math/ivp/_api.py | 17 +---------------- 1 file changed, 1 insertion(+), 16 deletions(-) diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_api.py index 5a9d0a2eb..2d93b7c4e 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_api.py @@ -297,11 +297,6 @@ def solve_ivp( if events is not None: events = [lambda t, x, event=event: event(t, x, argk) for event in events] g = [event(t0, y0) for event in events] - t_events = [[] for _ in range(len(events))] - y_events = [[] for _ in range(len(events))] - else: - t_events = None - y_events = None status = None while status is None: @@ -324,22 +319,12 @@ def solve_ivp( g_new = [event(t, y) for event in events] active_events = find_active_events(g, g_new, event_dir) if active_events.size > 0: - if sol is None: - sol = solver.dense_output() - - root_indices, roots, terminate = handle_events( + _, roots, terminate = handle_events( sol, events, active_events, is_terminal, t_old, t ) - - for e, te in zip(root_indices, roots): - t_events[e].append(te) - y_events[e].append(sol(te)) - if terminate: status = 1 t = roots[-1] - y = sol(t) - g = g_new ts.append(t) From 3d9aeb99c0b2ba553ac8978ea4cbf8b284f92bf9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 10:29:25 +0100 Subject: [PATCH 183/346] cleanup --- src/hapsira/core/math/ivp/__init__.py | 7 +++- src/hapsira/core/math/ivp/_const.py | 31 +++++++++++++++ src/hapsira/core/math/ivp/_rk.py | 38 +++++++++---------- src/hapsira/core/math/ivp/_rkerror.py | 2 +- src/hapsira/core/math/ivp/_rkstep.py | 13 +------ .../math/ivp/{_common.py => _solution.py} | 0 .../core/math/ivp/{_api.py => _solve.py} | 23 ++++++----- src/hapsira/core/math/linalg.py | 6 +++ 8 files changed, 75 insertions(+), 45 deletions(-) create mode 100644 src/hapsira/core/math/ivp/_const.py rename src/hapsira/core/math/ivp/{_common.py => _solution.py} (100%) rename src/hapsira/core/math/ivp/{_api.py => _solve.py} (96%) diff --git a/src/hapsira/core/math/ivp/__init__.py b/src/hapsira/core/math/ivp/__init__.py index 4c1c497c1..2aec5b6bf 100644 --- a/src/hapsira/core/math/ivp/__init__.py +++ b/src/hapsira/core/math/ivp/__init__.py @@ -1,4 +1,7 @@ -from ._api import solve_ivp +from ._solve import solve_ivp from ._brentq import brentq -__all__ = ["solve_ivp", "brentq"] +__all__ = [ + "solve_ivp", + "brentq", +] diff --git a/src/hapsira/core/math/ivp/_const.py b/src/hapsira/core/math/ivp/_const.py new file mode 100644 index 000000000..b885184d2 --- /dev/null +++ b/src/hapsira/core/math/ivp/_const.py @@ -0,0 +1,31 @@ +__all__ = [ + "N_RV", + "N_STAGES", + "SAFETY", + "MIN_FACTOR", + "MAX_FACTOR", + "INTERPOLATOR_POWER", + "N_STAGES_EXTENDED", + "ERROR_ESTIMATOR_ORDER", + "ERROR_EXPONENT", + "KSIG", +] + +N_RV = 6 +N_STAGES = 12 + +SAFETY = 0.9 # Multiply steps computed from asymptotic behaviour of errors by this. + +MIN_FACTOR = 0.2 # Minimum allowed decrease in a step size. +MAX_FACTOR = 10 # Maximum allowed increase in a step size. + +INTERPOLATOR_POWER = 7 +N_STAGES_EXTENDED = 16 +ERROR_ESTIMATOR_ORDER = 7 +ERROR_EXPONENT = -1 / (ERROR_ESTIMATOR_ORDER + 1) + +KSIG = ( + "Tuple([" + + ",".join(["Tuple([" + ",".join(["f"] * N_RV) + "])"] * (N_STAGES + 1)) + + "])" +) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 6a882dd65..cbeed0ba8 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -3,8 +3,20 @@ import numpy as np +from ._const import ( + N_RV, + N_STAGES, + KSIG, + SAFETY, + MIN_FACTOR, + MAX_FACTOR, + INTERPOLATOR_POWER, + N_STAGES_EXTENDED, + ERROR_ESTIMATOR_ORDER, + ERROR_EXPONENT, +) from ._dop853_coefficients import A as _A, C as _C, D as _D -from ._rkstep import rk_step_hf, N_RV, N_STAGES, KSIG +from ._rkstep import rk_step_hf from ._rkerror import estimate_error_norm_hf from ...jit import array_to_V_hf, hjit, DSIG @@ -16,6 +28,7 @@ max_VV_hf, mul_Vs_hf, nextafter_hf, + norm_VV_hf, sub_VV_hf, EPS, ) @@ -25,23 +38,6 @@ ] -# Multiply steps computed from asymptotic behaviour of errors by this. -SAFETY = 0.9 - -MIN_FACTOR = 0.2 # Minimum allowed decrease in a step size. -MAX_FACTOR = 10 # Maximum allowed increase in a step size. - -INTERPOLATOR_POWER = 7 -N_STAGES_EXTENDED = 16 -ERROR_ESTIMATOR_ORDER = 7 -ERROR_EXPONENT = -1 / (ERROR_ESTIMATOR_ORDER + 1) - - -@hjit("f(V,V)") -def _norm_VV_hf(x, y): - return sqrt(x[0] ** 2 + x[1] ** 2 + x[2] ** 2 + y[0] ** 2 + y[1] ** 2 + y[2] ** 2) - - @hjit(f"f(F({DSIG:s}),f,V,V,f,V,V,f,f,f,f)") def _select_initial_step_hf( fun, t0, rr, vv, argk, fr, fv, direction, order, rtol, atol @@ -58,8 +54,8 @@ def _select_initial_step_hf( ) factor = 1 / sqrt(6) - d0 = _norm_VV_hf(div_VV_hf(rr, scale_r), div_VV_hf(vv, scale_v)) * factor - d1 = _norm_VV_hf(div_VV_hf(fr, scale_r), div_VV_hf(fv, scale_v)) * factor + d0 = norm_VV_hf(div_VV_hf(rr, scale_r), div_VV_hf(vv, scale_v)) * factor + d1 = norm_VV_hf(div_VV_hf(fr, scale_r), div_VV_hf(fv, scale_v)) * factor if d0 < 1e-5 or d1 < 1e-5: h0 = 1e-6 @@ -77,7 +73,7 @@ def _select_initial_step_hf( ) d2 = ( - _norm_VV_hf( + norm_VV_hf( div_VV_hf(sub_VV_hf(fr1, fr), scale_r), div_VV_hf(sub_VV_hf(fv1, fv), scale_v), ) diff --git a/src/hapsira/core/math/ivp/_rkerror.py b/src/hapsira/core/math/ivp/_rkerror.py index 7a5a0a1cb..b8f1f20db 100644 --- a/src/hapsira/core/math/ivp/_rkerror.py +++ b/src/hapsira/core/math/ivp/_rkerror.py @@ -1,7 +1,7 @@ from math import sqrt +from ._const import N_RV, KSIG from ._dop853_coefficients import E3 as _E3, E5 as _E5 -from ._rkstep import N_RV, KSIG from ...jit import hjit from ....settings import settings diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index 83d5aa8f1..83d031811 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -1,3 +1,4 @@ +from ._const import N_STAGES, KSIG from ._dop853_coefficients import A as _A, B as _B, C as _C from ..linalg import add_VV_hf from ...jit import hjit, DSIG @@ -15,14 +16,8 @@ __all__ = [ "rk_step_hf", - "N_RV", - "N_STAGES", - "KSIG", ] -N_RV = 6 -N_STAGES = 12 - A01 = tuple(float_(number) for number in _A[1, :N_STAGES]) A02 = tuple(float_(number) for number in _A[2, :N_STAGES]) A03 = tuple(float_(number) for number in _A[3, :N_STAGES]) @@ -37,12 +32,6 @@ B = tuple(float_(number) for number in _B) C = tuple(float_(number) for number in _C[:N_STAGES]) -KSIG = ( - "Tuple([" - + ",".join(["Tuple([" + ",".join(["f"] * N_RV) + "])"] * (N_STAGES + 1)) - + "])" -) - @hjit(f"Tuple([V,V,V,V,{KSIG:s}])(F({DSIG:s}),f,V,V,V,V,f,f)") def rk_step_hf(fun, t, rr, vv, fr, fv, h, argk): diff --git a/src/hapsira/core/math/ivp/_common.py b/src/hapsira/core/math/ivp/_solution.py similarity index 100% rename from src/hapsira/core/math/ivp/_common.py rename to src/hapsira/core/math/ivp/_solution.py diff --git a/src/hapsira/core/math/ivp/_api.py b/src/hapsira/core/math/ivp/_solve.py similarity index 96% rename from src/hapsira/core/math/ivp/_api.py rename to src/hapsira/core/math/ivp/_solve.py index 2d93b7c4e..063ab281b 100644 --- a/src/hapsira/core/math/ivp/_api.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -3,12 +3,17 @@ import numpy as np from ._brentq import brentq -from ._common import OdeSolution +from ._solution import OdeSolution from ._rk import DOP853 from ...math.linalg import EPS -def solve_event_equation(event, sol, t_old, t): +__all__ = [ + "solve_ivp", +] + + +def _solve_event_equation(event, sol, t_old, t): """Solve an equation corresponding to an ODE event. The equation is ``event(t, y(t)) = 0``, here ``y(t)`` is known from an @@ -35,7 +40,7 @@ def solve_event_equation(event, sol, t_old, t): return brentq(lambda t: event(t, sol(t)), t_old, t, xtol=4 * EPS, rtol=4 * EPS) -def handle_events(sol, events, active_events, is_terminal, t_old, t): +def _handle_events(sol, events, active_events, is_terminal, t_old, t): """Helper function to handle events. Parameters @@ -63,7 +68,7 @@ def handle_events(sol, events, active_events, is_terminal, t_old, t): Whether a terminal event occurred. """ roots = [ - solve_event_equation(events[event_index], sol, t_old, t) + _solve_event_equation(events[event_index], sol, t_old, t) for event_index in active_events ] @@ -86,7 +91,7 @@ def handle_events(sol, events, active_events, is_terminal, t_old, t): return active_events, roots, terminate -def prepare_events(events): +def _prepare_events(events): """Standardize event functions and extract is_terminal and direction.""" if callable(events): events = (events,) @@ -111,7 +116,7 @@ def prepare_events(events): return events, is_terminal, direction -def find_active_events(g, g_new, direction): +def _find_active_events(g, g_new, direction): """Find which event occurred during an integration step. Parameters @@ -292,7 +297,7 @@ def solve_ivp( interpolants = [] - events, is_terminal, event_dir = prepare_events(events) + events, is_terminal, event_dir = _prepare_events(events) if events is not None: events = [lambda t, x, event=event: event(t, x, argk) for event in events] @@ -317,9 +322,9 @@ def solve_ivp( if events is not None: g_new = [event(t, y) for event in events] - active_events = find_active_events(g, g_new, event_dir) + active_events = _find_active_events(g, g_new, event_dir) if active_events.size > 0: - _, roots, terminate = handle_events( + _, roots, terminate = _handle_events( sol, events, active_events, is_terminal, t_old, t ) if terminate: diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 4c2e78087..1cc42e171 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -20,6 +20,7 @@ "nextafter_hf", "norm_hf", "norm_vf", + "norm_VV_hf", "sign_hf", "sub_VV_hf", "transpose_M_hf", @@ -156,6 +157,11 @@ def norm_vf(a, b, c): return norm_hf((a, b, c)) +@hjit("f(V,V)") +def norm_VV_hf(x, y): + return sqrt(x[0] ** 2 + x[1] ** 2 + x[2] ** 2 + y[0] ** 2 + y[1] ** 2 + y[2] ** 2) + + @hjit("f(f)") def sign_hf(x): if x < 0.0: From 466867a36d50166f0f43831e9fe6947348236982 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 10:31:34 +0100 Subject: [PATCH 184/346] fix const --- src/hapsira/core/math/ivp/_dop853_coefficients.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/ivp/_dop853_coefficients.py b/src/hapsira/core/math/ivp/_dop853_coefficients.py index 7f9d03aff..a425fc8ff 100644 --- a/src/hapsira/core/math/ivp/_dop853_coefficients.py +++ b/src/hapsira/core/math/ivp/_dop853_coefficients.py @@ -1,8 +1,15 @@ import numpy as np -N_STAGES = 12 -N_STAGES_EXTENDED = 16 -INTERPOLATOR_POWER = 7 +from ._const import N_STAGES, N_STAGES_EXTENDED, INTERPOLATOR_POWER + +__all__ = [ + "A", + "B", + "C", + "D", + "E3", + "E5", +] C = np.array( [ From b6eba2d40a1c6f3ac1aad4fcf14ef369acec5aec Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 10:32:24 +0100 Subject: [PATCH 185/346] fix const --- src/hapsira/core/math/ivp/_const.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/hapsira/core/math/ivp/_const.py b/src/hapsira/core/math/ivp/_const.py index b885184d2..7ed0b3bc4 100644 --- a/src/hapsira/core/math/ivp/_const.py +++ b/src/hapsira/core/math/ivp/_const.py @@ -13,6 +13,7 @@ N_RV = 6 N_STAGES = 12 +N_STAGES_EXTENDED = 16 SAFETY = 0.9 # Multiply steps computed from asymptotic behaviour of errors by this. From 4613d1b1da6e8e52cda66b4ed3a82f538d2a4826 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 10:49:43 +0100 Subject: [PATCH 186/346] rename norm func --- src/hapsira/core/czml_utils.py | 6 ++--- src/hapsira/core/elements.py | 24 +++++++++++-------- src/hapsira/core/events.py | 8 +++---- src/hapsira/core/flybys.py | 8 +++---- src/hapsira/core/iod.py | 14 +++++------ src/hapsira/core/maneuver.py | 16 ++++++------- src/hapsira/core/math/ivp/_rk.py | 4 ++-- src/hapsira/core/math/ivp/_rkerror.py | 4 ++-- src/hapsira/core/math/linalg.py | 6 ++--- src/hapsira/core/perturbations.py | 20 ++++++++-------- src/hapsira/core/propagation/vallado.py | 4 ++-- src/hapsira/core/spheroid_location.py | 10 ++++---- src/hapsira/core/thrust/change_a_inc.py | 13 +++++++--- src/hapsira/core/thrust/change_argp.py | 13 +++++++--- src/hapsira/core/thrust/change_ecc_inc.py | 15 ++++++++---- .../core/thrust/change_ecc_quasioptimal.py | 6 ++--- tests/tests_twobody/test_perturbations.py | 10 ++++---- tests/tests_twobody/test_propagation.py | 4 ++-- 18 files changed, 105 insertions(+), 80 deletions(-) diff --git a/src/hapsira/core/czml_utils.py b/src/hapsira/core/czml_utils.py index eb3b275da..4e7a515e1 100644 --- a/src/hapsira/core/czml_utils.py +++ b/src/hapsira/core/czml_utils.py @@ -1,7 +1,7 @@ from numba import njit as jit import numpy as np -from .math.linalg import norm_hf +from .math.linalg import norm_V_hf @jit @@ -95,7 +95,7 @@ def project_point_on_ellipsoid(x, y, z, a, b, c): """ p1, p2 = intersection_ellipsoid_line(x, y, z, x, y, z, a, b, c) - norm_1 = norm_hf((p1[0] - x, p1[1] - y, p1[2] - z)) - norm_2 = norm_hf((p2[0] - x, p2[1] - y, p2[2] - z)) + norm_1 = norm_V_hf((p1[0] - x, p1[1] - y, p1[2] - z)) + norm_2 = norm_V_hf((p2[0] - x, p2[1] - y, p2[2] - z)) return p1 if norm_1 <= norm_2 else p2 diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index 0ec7c0e8b..f9a948451 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -13,7 +13,7 @@ matmul_VM_hf, matmul_VV_hf, mul_Vs_hf, - norm_hf, + norm_V_hf, sub_VV_hf, transpose_M_hf, ) @@ -63,7 +63,7 @@ def eccentricity_vector_hf(k, r, v): v : tuple[float,float,float] Velocity vector (km / s) """ - a = matmul_VV_hf(v, v) - k / norm_hf(r) + a = matmul_VV_hf(v, v) - k / norm_V_hf(r) b = matmul_VV_hf(r, v) return div_Vs_hf(sub_VV_hf(mul_Vs_hf(r, a), mul_Vs_hf(v, b)), k) @@ -448,14 +448,14 @@ def rv2coe_hf(k, r, v, tol): n = cross_VV_hf((0, 0, 1), h) e = mul_Vs_hf( sub_VV_hf( - mul_Vs_hf(r, (matmul_VV_hf(v, v) - k / norm_hf(r))), + mul_Vs_hf(r, (matmul_VV_hf(v, v) - k / norm_V_hf(r))), mul_Vs_hf(v, matmul_VV_hf(r, v)), ), 1 / k, ) - ecc = norm_hf(e) + ecc = norm_V_hf(e) p = matmul_VV_hf(h, h) / k - inc = acos(h[2] / norm_hf(h)) + inc = acos(h[2] / norm_V_hf(h)) circular = ecc < tol equatorial = abs(inc) < tol @@ -463,12 +463,16 @@ def rv2coe_hf(k, r, v, tol): if equatorial and not circular: raan = 0 argp = atan2(e[1], e[0]) % (2 * pi) # Longitude of periapsis - nu = atan2(matmul_VV_hf(h, cross_VV_hf(e, r)) / norm_hf(h), matmul_VV_hf(r, e)) + nu = atan2( + matmul_VV_hf(h, cross_VV_hf(e, r)) / norm_V_hf(h), matmul_VV_hf(r, e) + ) elif not equatorial and circular: raan = atan2(n[1], n[0]) % (2 * pi) argp = 0 # Argument of latitude - nu = atan2(matmul_VV_hf(r, cross_VV_hf(h, n)) / norm_hf(h), matmul_VV_hf(r, n)) + nu = atan2( + matmul_VV_hf(r, cross_VV_hf(h, n)) / norm_V_hf(h), matmul_VV_hf(r, n) + ) elif equatorial and circular: raan = 0 argp = 0 @@ -478,16 +482,16 @@ def rv2coe_hf(k, r, v, tol): ka = k * a if a > 0: e_se = matmul_VV_hf(r, v) / sqrt(ka) - e_ce = norm_hf(r) * matmul_VV_hf(v, v) / k - 1 + e_ce = norm_V_hf(r) * matmul_VV_hf(v, v) / k - 1 nu = E_to_nu_hf(atan2(e_se, e_ce), ecc) else: e_sh = matmul_VV_hf(r, v) / sqrt(-ka) - e_ch = norm_hf(r) * (norm_hf(v) ** 2) / k - 1 + e_ch = norm_V_hf(r) * (norm_V_hf(v) ** 2) / k - 1 nu = F_to_nu_hf(log((e_ch + e_sh) / (e_ch - e_sh)) / 2, ecc) raan = atan2(n[1], n[0]) % (2 * pi) px = matmul_VV_hf(r, n) - py = matmul_VV_hf(r, cross_VV_hf(h, n)) / norm_hf(h) + py = matmul_VV_hf(r, cross_VV_hf(h, n)) / norm_V_hf(h) argp = (atan2(py, px) - nu) % (2 * pi) nu = (nu + pi) % (2 * pi) - pi diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index a0a049acb..3b6d7148b 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -5,7 +5,7 @@ from .elements import coe_rotation_matrix_hf, rv2coe_hf, RV2COE_TOL from .jit import array_to_V_hf, hjit, gjit -from .math.linalg import norm_hf, matmul_VV_hf +from .math.linalg import norm_V_hf, matmul_VV_hf from .util import planetocentric_to_AltAz_hf @@ -53,7 +53,7 @@ def eclipse_function(k, u_, r_sec, R_sec, R_primary, umbra=True): # Make arrays contiguous for faster dot product with numba. P_, Q_ = np.ascontiguousarray(PQW[:, 0]), np.ascontiguousarray(PQW[:, 1]) - r_sec_norm = norm_hf(array_to_V_hf(r_sec)) + r_sec_norm = norm_V_hf(array_to_V_hf(r_sec)) beta = (P_ @ r_sec) / r_sec_norm zeta = (Q_ @ r_sec) / r_sec_norm @@ -90,8 +90,8 @@ def line_of_sight_hf(r1, r2, R): located by r1 and r2, else negative. """ - r1_norm = norm_hf(r1) - r2_norm = norm_hf(r2) + r1_norm = norm_V_hf(r1) + r2_norm = norm_V_hf(r2) theta = acos(matmul_VV_hf(r1, r2) / r1_norm / r2_norm) theta_1 = acos(R / r1_norm) diff --git a/src/hapsira/core/flybys.py b/src/hapsira/core/flybys.py index 18064ff0a..53339ff60 100644 --- a/src/hapsira/core/flybys.py +++ b/src/hapsira/core/flybys.py @@ -5,7 +5,7 @@ from numpy import cross from .jit import array_to_V_hf -from .math.linalg import norm_hf +from .math.linalg import norm_V_hf @jit @@ -34,7 +34,7 @@ def compute_flyby(v_spacecraft, v_body, k, r_p, theta): """ v_inf_1 = v_spacecraft - v_body # Hyperbolic excess velocity - v_inf = norm_hf(array_to_V_hf(v_inf_1)) + v_inf = norm_V_hf(array_to_V_hf(v_inf_1)) ecc = 1 + r_p * v_inf**2 / k # Eccentricity of the entry hyperbola delta = 2 * np.arcsin(1 / ecc) # Turn angle @@ -49,7 +49,7 @@ def compute_flyby(v_spacecraft, v_body, k, r_p, theta): S_vec = v_inf_1 / v_inf c_vec = np.array([0, 0, 1]) T_vec = cross(S_vec, c_vec) - T_vec = T_vec / norm_hf(array_to_V_hf(T_vec)) + T_vec = T_vec / norm_V_hf(array_to_V_hf(T_vec)) R_vec = cross(S_vec, T_vec) # This vector defines the B-Plane @@ -63,7 +63,7 @@ def compute_flyby(v_spacecraft, v_body, k, r_p, theta): # And now we rotate the outbound hyperbolic excess velocity # u_vec = v_inf_1 / norm(v_inf) = S_vec v_vec = cross(rot_v, v_inf_1) - v_vec = v_vec / norm_hf(array_to_V_hf(v_vec)) + v_vec = v_vec / norm_V_hf(array_to_V_hf(v_vec)) v_inf_2 = v_inf * (np.cos(delta) * S_vec + np.sin(delta) * v_vec) diff --git a/src/hapsira/core/iod.py b/src/hapsira/core/iod.py index a17bb2487..d7cf19793 100644 --- a/src/hapsira/core/iod.py +++ b/src/hapsira/core/iod.py @@ -8,7 +8,7 @@ div_Vs_hf, matmul_VV_hf, mul_Vs_hf, - norm_hf, + norm_V_hf, sub_VV_hf, ) from .math.special import hyp2f1b_hf, stumpff_c2_hf, stumpff_c3_hf @@ -396,8 +396,8 @@ def vallado_hf(k, r0, r, tof, M, prograde, lowpath, numiter, rtol): t_m = 1 if prograde else -1 - norm_r0 = norm_hf(r0) - norm_r = norm_hf(r) + norm_r0 = norm_V_hf(r0) + norm_r = norm_V_hf(r) norm_r0_times_norm_r = norm_r0 * norm_r norm_r0_plus_norm_r = norm_r0 + norm_r @@ -521,9 +521,9 @@ def izzo_hf(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol): # Chord c = sub_VV_hf(r2, r1) c_norm, r1_norm, r2_norm = ( - norm_hf(c), - norm_hf(r1), - norm_hf(r2), + norm_V_hf(c), + norm_V_hf(r1), + norm_V_hf(r2), ) # Semiperimeter @@ -532,7 +532,7 @@ def izzo_hf(k, r1, r2, tof, M, prograde, lowpath, numiter, rtol): # Versors i_r1, i_r2 = div_Vs_hf(r1, r1_norm), div_Vs_hf(r2, r2_norm) i_h = cross_VV_hf(i_r1, i_r2) - i_h = div_Vs_hf(i_h, norm_hf(i_h)) # Fixed from paper + i_h = div_Vs_hf(i_h, norm_V_hf(i_h)) # Fixed from paper # Geometry of the problem ll = sqrt(1 - min(1.0, c_norm / s)) diff --git a/src/hapsira/core/maneuver.py b/src/hapsira/core/maneuver.py index 2b6800331..25559aaa4 100644 --- a/src/hapsira/core/maneuver.py +++ b/src/hapsira/core/maneuver.py @@ -12,7 +12,7 @@ ) from .jit import array_to_V_hf -from .math.linalg import norm_hf +from .math.linalg import norm_V_hf @jit @@ -50,13 +50,13 @@ def hohmann(k, rv, r_f): _, ecc, inc, raan, argp, nu = rv2coe_hf( k, array_to_V_hf(rv[0]), array_to_V_hf(rv[1]), RV2COE_TOL ) - h_i = norm_hf(array_to_V_hf(cross(*rv))) + h_i = norm_V_hf(array_to_V_hf(cross(*rv))) p_i = h_i**2 / k r_i, v_i = rv_pqw_hf(k, p_i, ecc, nu) - r_i = norm_hf(r_i) - v_i = norm_hf(v_i) + r_i = norm_V_hf(r_i) + v_i = norm_V_hf(v_i) a_trans = (r_i + r_f) / 2 dv_a = np.sqrt(2 * k / r_i - k / a_trans) - v_i @@ -122,13 +122,13 @@ def bielliptic(k, r_b, r_f, rv): _, ecc, inc, raan, argp, nu = rv2coe_hf( k, array_to_V_hf(rv[0]), array_to_V_hf(rv[1]), RV2COE_TOL ) - h_i = norm_hf(array_to_V_hf(cross(*rv))) + h_i = norm_V_hf(array_to_V_hf(cross(*rv))) p_i = h_i**2 / k r_i, v_i = rv_pqw_hf(k, p_i, ecc, nu) - r_i = norm_hf(r_i) - v_i = norm_hf(v_i) + r_i = norm_V_hf(r_i) + v_i = norm_V_hf(v_i) a_trans1 = (r_i + r_b) / 2 a_trans2 = (r_b + r_f) / 2 @@ -201,6 +201,6 @@ def correct_pericenter(k, R, J2, max_delta_r, v, a, inc, ecc): delta_t = abs(delta_w / dw) delta_v = 0.5 * n * a * ecc * abs(delta_w) - vf_ = v / norm_hf(array_to_V_hf(v)) * delta_v + vf_ = v / norm_V_hf(array_to_V_hf(v)) * delta_v return delta_t, vf_ diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index cbeed0ba8..3c235be2e 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -17,7 +17,7 @@ ) from ._dop853_coefficients import A as _A, C as _C, D as _D from ._rkstep import rk_step_hf -from ._rkerror import estimate_error_norm_hf +from ._rkerror import estimate_error_norm_V_hf from ...jit import array_to_V_hf, hjit, DSIG from ...math.linalg import ( @@ -159,7 +159,7 @@ def _step_impl_hf( ), atol, ) - error_norm = estimate_error_norm_hf( + error_norm = estimate_error_norm_V_hf( K_new, h, scale_r, diff --git a/src/hapsira/core/math/ivp/_rkerror.py b/src/hapsira/core/math/ivp/_rkerror.py index b8f1f20db..a4eb1d152 100644 --- a/src/hapsira/core/math/ivp/_rkerror.py +++ b/src/hapsira/core/math/ivp/_rkerror.py @@ -16,7 +16,7 @@ __all__ = [ - "estimate_error_norm_hf", + "estimate_error_norm_V_hf", ] @@ -25,7 +25,7 @@ @hjit(f"f({KSIG:s},f,V,V)") -def estimate_error_norm_hf(K, h, scale_r, scale_v): +def estimate_error_norm_V_hf(K, h, scale_r, scale_v): K00, K01, K02, K03, K04, K05, K06, K07, K08, K09, K10, K11, K12 = K err3 = ( diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 1cc42e171..0c3147fdb 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -18,7 +18,7 @@ "mul_Vs_hf", "mul_VV_hf", "nextafter_hf", - "norm_hf", + "norm_V_hf", "norm_vf", "norm_VV_hf", "sign_hf", @@ -147,14 +147,14 @@ def nextafter_hf(x, direction): @hjit("f(V)") -def norm_hf(a): +def norm_V_hf(a): return sqrt(matmul_VV_hf(a, a)) @vjit("f(f,f,f)") def norm_vf(a, b, c): # TODO add axis setting in some way for util.norm? - return norm_hf((a, b, c)) + return norm_V_hf((a, b, c)) @hjit("f(V,V)") diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index ae9ecbbee..15e42aac8 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -2,7 +2,7 @@ from .events import line_of_sight_hf from .jit import hjit -from .math.linalg import norm_hf, mul_Vs_hf, mul_VV_hf, sub_VV_hf +from .math.linalg import norm_V_hf, mul_Vs_hf, mul_VV_hf, sub_VV_hf __all__ = [ @@ -46,7 +46,7 @@ def J2_perturbation_hf(t0, rr, vv, k, J2, R): Howard Curtis, (12.30) """ - r = norm_hf(rr) + r = norm_V_hf(rr) factor = (3.0 / 2.0) * k * J2 * (R**2) / (r**5) @@ -81,7 +81,7 @@ def J3_perturbation_hf(t0, rr, vv, k, J3, R): This perturbation has not been fully validated, see https://github.com/poliastro/poliastro/pull/398 """ - r = norm_hf(rr) + r = norm_V_hf(rr) factor = (1.0 / 2.0) * k * J3 * (R**3) / (r**5) cos_phi = rr[2] / r @@ -131,9 +131,9 @@ def atmospheric_drag_exponential_hf(t0, rr, vv, k, R, C_D, A_over_m, H0, rho0): the atmospheric density model is rho(H) = rho0 x exp(-H / H0) """ - H = norm_hf(rr) + H = norm_V_hf(rr) - v = norm_hf(vv) + v = norm_V_hf(vv) B = C_D * A_over_m rho = rho0 * exp(-(H - R) / H0) @@ -173,7 +173,7 @@ def atmospheric_drag_hf(t0, rr, vv, k, C_D, A_over_m, rho): computed by a model from hapsira.earth.atmosphere """ - v = norm_hf(vv) + v = norm_V_hf(vv) B = C_D * A_over_m return mul_Vs_hf(vv, -(1.0 / 2.0) * rho * B * v) @@ -212,8 +212,8 @@ def third_body_hf(t0, rr, vv, k, k_third, perturbation_body): body_r = perturbation_body(t0) delta_r = sub_VV_hf(body_r, rr) return sub_VV_hf( - mul_Vs_hf(delta_r, k_third / norm_hf(delta_r) ** 3), - mul_Vs_hf(body_r, k_third / norm_hf(body_r) ** 3), + mul_Vs_hf(delta_r, k_third / norm_V_hf(delta_r) ** 3), + mul_Vs_hf(body_r, k_third / norm_V_hf(body_r) ** 3), ) @@ -254,10 +254,10 @@ def radiation_pressure_hf(t0, rr, vv, k, R, C_R, A_over_m, Wdivc_s, star): """ r_star = star(t0) - P_s = Wdivc_s / (norm_hf(r_star) ** 2) + P_s = Wdivc_s / (norm_V_hf(r_star) ** 2) if line_of_sight_hf(rr, r_star, R) > 0: nu = 1.0 else: nu = 0.0 - return mul_Vs_hf(r_star, -nu * P_s * (C_R * A_over_m) / norm_hf(r_star)) + return mul_Vs_hf(r_star, -nu * P_s * (C_R * A_over_m) / norm_V_hf(r_star)) diff --git a/src/hapsira/core/propagation/vallado.py b/src/hapsira/core/propagation/vallado.py index 076892e31..df6ee56c7 100644 --- a/src/hapsira/core/propagation/vallado.py +++ b/src/hapsira/core/propagation/vallado.py @@ -1,7 +1,7 @@ from math import log, sqrt from ..elements import coe2rv_hf, rv2coe_hf, RV2COE_TOL -from ..math.linalg import add_VV_hf, matmul_VV_hf, mul_Vs_hf, norm_hf, sign_hf +from ..math.linalg import add_VV_hf, matmul_VV_hf, mul_Vs_hf, norm_V_hf, sign_hf from ..math.special import stumpff_c2_hf, stumpff_c3_hf from ..jit import array_to_V_hf, hjit, vjit, gjit @@ -86,7 +86,7 @@ def _vallado_hf(k, r0, v0, tof, numiter): """ # Cache some results dot_r0v0 = matmul_VV_hf(r0, v0) - norm_r0 = norm_hf(r0) + norm_r0 = norm_V_hf(r0) sqrt_mu = k**0.5 alpha = -matmul_VV_hf(v0, v0) / k + 2 / norm_r0 diff --git a/src/hapsira/core/spheroid_location.py b/src/hapsira/core/spheroid_location.py index 86fe567fc..c92ff1095 100644 --- a/src/hapsira/core/spheroid_location.py +++ b/src/hapsira/core/spheroid_location.py @@ -4,14 +4,14 @@ import numpy as np from .jit import array_to_V_hf -from .math.linalg import norm_hf +from .math.linalg import norm_V_hf @jit def cartesian_cords(a, c, lon, lat, h): """Calculates cartesian coordinates. - Parametersnorm_hf + Parametersnorm_V_hf ---------- a : float Semi-major axis @@ -67,7 +67,7 @@ def N(a, b, c, cartesian_cords): """ x, y, z = cartesian_cords N = np.array([2 * x / a**2, 2 * y / b**2, 2 * z / c**2]) - N /= norm_hf(array_to_V_hf(N)) + N /= norm_V_hf(array_to_V_hf(N)) return N @@ -83,7 +83,7 @@ def tangential_vecs(N): """ u = np.array([1.0, 0, 0]) u -= (u @ N) * N - u /= norm_hf(array_to_V_hf(u)) + u /= norm_V_hf(array_to_V_hf(u)) v = np.cross(N, u) return u, v @@ -126,7 +126,7 @@ def distance(cartesian_cords, px, py, pz): """ c = cartesian_cords u = np.array([px, py, pz]) - d = norm_hf(array_to_V_hf(c - u)) + d = norm_V_hf(array_to_V_hf(c - u)) return d diff --git a/src/hapsira/core/thrust/change_a_inc.py b/src/hapsira/core/thrust/change_a_inc.py index 499887c98..bf4f3b9b5 100644 --- a/src/hapsira/core/thrust/change_a_inc.py +++ b/src/hapsira/core/thrust/change_a_inc.py @@ -2,7 +2,14 @@ from ..jit import hjit, gjit from ..elements import circular_velocity_hf -from ..math.linalg import add_VV_hf, cross_VV_hf, div_Vs_hf, mul_Vs_hf, norm_hf, sign_hf +from ..math.linalg import ( + add_VV_hf, + cross_VV_hf, + div_Vs_hf, + mul_Vs_hf, + norm_V_hf, + sign_hf, +) __all__ = [ @@ -119,9 +126,9 @@ def a_d_hf(t0, rr, vv, k): # Change sign of beta with the out-of-plane velocity beta_ = _beta_hf(t0, V_0, f, beta_0_) * sign_hf(rr[0] * (inc_f - inc_0)) - t_ = div_Vs_hf(vv, norm_hf(vv)) + t_ = div_Vs_hf(vv, norm_V_hf(vv)) crv = cross_VV_hf(rr, vv) - w_ = div_Vs_hf(crv, norm_hf(crv)) + w_ = div_Vs_hf(crv, norm_V_hf(crv)) accel_v = mul_Vs_hf( add_VV_hf(mul_Vs_hf(t_, cos(beta_)), mul_Vs_hf(w_, sin(beta_))), f ) diff --git a/src/hapsira/core/thrust/change_argp.py b/src/hapsira/core/thrust/change_argp.py index 1ba715b2a..ab2d28e55 100644 --- a/src/hapsira/core/thrust/change_argp.py +++ b/src/hapsira/core/thrust/change_argp.py @@ -2,7 +2,14 @@ from ..elements import circular_velocity_hf, rv2coe_hf, RV2COE_TOL from ..jit import hjit, gjit -from ..math.linalg import add_VV_hf, cross_VV_hf, div_Vs_hf, mul_Vs_hf, norm_hf, sign_hf +from ..math.linalg import ( + add_VV_hf, + cross_VV_hf, + div_Vs_hf, + mul_Vs_hf, + norm_V_hf, + sign_hf, +) __all__ = [ @@ -71,9 +78,9 @@ def a_d_hf(t0, rr, vv, k): alpha_ = nu - pi / 2 - r_ = div_Vs_hf(rr, norm_hf(rr)) + r_ = div_Vs_hf(rr, norm_V_hf(rr)) crv = cross_VV_hf(rr, vv) - w_ = div_Vs_hf(crv, norm_hf(crv)) + w_ = div_Vs_hf(crv, norm_V_hf(crv)) s_ = cross_VV_hf(w_, r_) accel_v = mul_Vs_hf( add_VV_hf(mul_Vs_hf(s_, cos(alpha_)), mul_Vs_hf(r_, sin(alpha_))), f diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index 0cff79351..2964046b0 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -17,7 +17,14 @@ RV2COE_TOL, ) from ..jit import array_to_V_hf, hjit, vjit, gjit -from ..math.linalg import add_VV_hf, cross_VV_hf, div_Vs_hf, mul_Vs_hf, norm_hf, sign_hf +from ..math.linalg import ( + add_VV_hf, + cross_VV_hf, + div_Vs_hf, + mul_Vs_hf, + norm_V_hf, + sign_hf, +) __all__ = [ @@ -91,10 +98,10 @@ def _prepare_hf(k, a, ecc_0, ecc_f, inc_0, inc_f, argp, r, v, f): e_vec = eccentricity_vector_hf(k, r, v) ref_vec = div_Vs_hf(e_vec, ecc_0) else: - ref_vec = div_Vs_hf(r, norm_hf(r)) + ref_vec = div_Vs_hf(r, norm_V_hf(r)) h_vec = cross_VV_hf(r, v) # Specific angular momentum vector - h_unit = div_Vs_hf(h_vec, norm_hf(h_vec)) + h_unit = div_Vs_hf(h_vec, norm_V_hf(h_vec)) thrust_unit = mul_Vs_hf(cross_VV_hf(h_unit, ref_vec), sign_hf(ecc_f - ecc_0)) beta_0 = beta_hf(ecc_0, ecc_f, inc_0, inc_f, argp) @@ -135,7 +142,7 @@ def a_d_hf(t0, rr, vv, k_): ) # The sign of ß reverses at minor axis crossings w_ = mul_Vs_hf( - cross_VV_hf(rr, vv), sign_hf(inc_f - inc_0) / norm_hf(cross_VV_hf(rr, vv)) + cross_VV_hf(rr, vv), sign_hf(inc_f - inc_0) / norm_V_hf(cross_VV_hf(rr, vv)) ) accel_v = mul_Vs_hf( add_VV_hf(mul_Vs_hf(thrust_unit, cos(beta_)), mul_Vs_hf(w_, sin(beta_))), f diff --git a/src/hapsira/core/thrust/change_ecc_quasioptimal.py b/src/hapsira/core/thrust/change_ecc_quasioptimal.py index 582ec9217..fb1ac52d4 100644 --- a/src/hapsira/core/thrust/change_ecc_quasioptimal.py +++ b/src/hapsira/core/thrust/change_ecc_quasioptimal.py @@ -4,7 +4,7 @@ from ..elements import circular_velocity_hf from ..jit import array_to_V_hf, hjit, gjit -from ..math.linalg import cross_VV_hf, div_Vs_hf, mul_Vs_hf, norm_hf, sign_hf +from ..math.linalg import cross_VV_hf, div_Vs_hf, mul_Vs_hf, norm_V_hf, sign_hf __all__ = [ "change_ecc_quasioptimal_hb", @@ -51,9 +51,9 @@ def _prepare_hf(k, a, ecc_0, ecc_f, e_vec, h_vec, r): if ecc_0 > 0.001: # Arbitrary tolerance ref_vec = div_Vs_hf(e_vec, ecc_0) else: - ref_vec = div_Vs_hf(r, norm_hf(r)) + ref_vec = div_Vs_hf(r, norm_V_hf(r)) - h_unit = div_Vs_hf(h_vec, norm_hf(h_vec)) + h_unit = div_Vs_hf(h_vec, norm_V_hf(h_vec)) thrust_unit = mul_Vs_hf(cross_VV_hf(h_unit, ref_vec), sign_hf(ecc_f - ecc_0)) return thrust_unit diff --git a/tests/tests_twobody/test_perturbations.py b/tests/tests_twobody/test_perturbations.py index 0cd525f45..99b7ecc62 100644 --- a/tests/tests_twobody/test_perturbations.py +++ b/tests/tests_twobody/test_perturbations.py @@ -11,7 +11,7 @@ from hapsira.core.earth.atmosphere.coesa76 import density_hf as coesa76_density_hf from hapsira.core.elements import rv2coe_gf, RV2COE_TOL from hapsira.core.jit import hjit, djit -from hapsira.core.math.linalg import add_VV_hf, mul_Vs_hf, norm_hf +from hapsira.core.math.linalg import add_VV_hf, mul_Vs_hf, norm_V_hf from hapsira.core.perturbations import ( # pylint: disable=E1120,E1136 J2_perturbation_hf, J3_perturbation_hf, @@ -385,7 +385,7 @@ def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) # Avoid undershooting H below attractor radius R - H = norm_hf(rr) + H = norm_V_hf(rr) if H < R: H = R @@ -440,7 +440,7 @@ def test_cowell_works_with_small_perturbations(): @hjit("V(f,V,V,f)") def accel_hf(t0, rr, vv, k): - return mul_Vs_hf(vv, 1e-5 / norm_hf(vv)) + return mul_Vs_hf(vv, 1e-5 / norm_V_hf(vv)) @djit def f_hf(t0, rr, vv, k): @@ -465,7 +465,7 @@ def test_cowell_converges_with_small_perturbations(): @hjit("V(f,V,V,f)") def accel_hf(t0, rr, vv, k): - norm_v = norm_hf(vv) + norm_v = norm_V_hf(vv) return mul_Vs_hf(vv, 0.0 / norm_v) @djit @@ -710,7 +710,7 @@ def test_solar_pressure(t_days, deltas_expected, sun_r): @hjit("V(f)") def sun_normalized_hf(t0): r = sun_r(t0) # sun_r is hf, returns V - return mul_Vs_hf(r, 149600000 / norm_hf(r)) + return mul_Vs_hf(r, 149600000 / norm_V_hf(r)) R_ = Earth.R.to(u.km).value Wdivc_s = Wdivc_sun.value diff --git a/tests/tests_twobody/test_propagation.py b/tests/tests_twobody/test_propagation.py index cd3dfd5e1..d57d302b3 100644 --- a/tests/tests_twobody/test_propagation.py +++ b/tests/tests_twobody/test_propagation.py @@ -11,7 +11,7 @@ from hapsira.constants import J2000 from hapsira.core.elements import rv2coe_gf, RV2COE_TOL from hapsira.core.jit import djit, hjit -from hapsira.core.math.linalg import add_VV_hf, mul_Vs_hf, norm_hf +from hapsira.core.math.linalg import add_VV_hf, mul_Vs_hf, norm_V_hf from hapsira.core.propagation.base import func_twobody_hf from hapsira.examples import iss from hapsira.frames import Planes @@ -293,7 +293,7 @@ def test_cowell_propagation_circle_to_circle(): @hjit("V(f,V,V,f)") def constant_accel_hf(t0, rr, vv, k): - norm_v = norm_hf(vv) + norm_v = norm_V_hf(vv) return mul_Vs_hf(vv, accel / norm_v) @djit From efdaeb07c813e84b91f6176b3ad40c91f0cbc093 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 10:55:36 +0100 Subject: [PATCH 187/346] rename norm vec func --- src/hapsira/core/math/linalg.py | 4 ++-- src/hapsira/twobody/events.py | 8 ++++---- src/hapsira/util.py | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 0c3147fdb..58f962862 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -19,7 +19,7 @@ "mul_VV_hf", "nextafter_hf", "norm_V_hf", - "norm_vf", + "norm_V_vf", "norm_VV_hf", "sign_hf", "sub_VV_hf", @@ -152,7 +152,7 @@ def norm_V_hf(a): @vjit("f(f,f,f)") -def norm_vf(a, b, c): +def norm_V_vf(a, b, c): # TODO add axis setting in some way for util.norm? return norm_V_hf((a, b, c)) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index dc931bd5b..2755c4df0 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -4,7 +4,7 @@ from astropy.coordinates import get_body_barycentric_posvel import numpy as np -from hapsira.core.math.linalg import norm_vf +from hapsira.core.math.linalg import norm_V_vf from hapsira.core.events import ( eclipse_function as eclipse_function_fast, line_of_sight_gf, @@ -71,7 +71,7 @@ def __init__(self, alt, R, terminal=True, direction=-1): def __call__(self, t, u, k): self._last_t = t - r_norm = norm_vf(*u[:3]) + r_norm = norm_V_vf(*u[:3]) return ( r_norm - self._R - self._alt @@ -121,7 +121,7 @@ def __init__(self, orbit, lat, terminal=False, direction=0): def __call__(self, t, u_, k): self._last_t = t - pos_on_body = (u_[:3] / norm_vf(*u_[:3])) * self._R + pos_on_body = (u_[:3] / norm_V_vf(*u_[:3])) * self._R _, lat_, _ = cartesian_to_ellipsoidal_fast(self._R, self._R_polar, *pos_on_body) return np.rad2deg(lat_) - self._lat @@ -273,7 +273,7 @@ def __init__(self, attractor, pos_coords, terminal=False, direction=0): def __call__(self, t, u_, k): self._last_t = t - if norm_vf(*u_[:3]) < self._R: + if norm_V_vf(*u_[:3]) < self._R: warn( "The norm of the position vector of the primary body is less than the radius of the attractor." ) diff --git a/src/hapsira/util.py b/src/hapsira/util.py index 074e0bad5..99b702ecf 100644 --- a/src/hapsira/util.py +++ b/src/hapsira/util.py @@ -4,7 +4,7 @@ from astropy.time import Time import numpy as np -from hapsira.core.math.linalg import norm_vf +from hapsira.core.math.linalg import norm_V_vf from hapsira.core.util import alinspace as alinspace_fast @@ -26,7 +26,7 @@ def norm(vec, axis=None): result = norm_np(vec.value, axis=axis) else: - result = norm_vf(*vec.value) + result = norm_V_vf(*vec.value) return result << vec.unit From a38e96832bdeee93ec47304f20535651f4540936 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 11:01:30 +0100 Subject: [PATCH 188/346] allow args in gjit --- src/hapsira/core/jit.py | 28 ++++++++++++++++++++++++---- 1 file changed, 24 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 79d702e6b..8fb6fd23e 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -135,14 +135,34 @@ def wrapper(inner_func: Callable) -> Callable: return wrapper -def djit(func): +def djit(*args, **kwargs) -> Callable: """ Wrapper for hjit to track differential equations """ - compiled = hjit(DSIG)(func) - compiled.djit = None # for debugging - return compiled + if len(args) == 1 and callable(args[0]): + outer_func = args[0] + args = tuple() + else: + outer_func = None + + def wrapper(inner_func: Callable) -> Callable: + """ + Applies JIT + """ + + compiled = hjit( + DSIG, + *args, + **kwargs, + )(inner_func) + compiled.djit = None # attribute for debugging + return compiled + + if outer_func is not None: + return wrapper(outer_func) + + return wrapper def vjit(*args, **kwargs) -> Callable: From 4b9e5bf002b57a243ff9dd44ac7de6705380639a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 11:04:09 +0100 Subject: [PATCH 189/346] add logging to djit --- src/hapsira/core/jit.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 8fb6fd23e..21ee45ca2 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -151,6 +151,13 @@ def wrapper(inner_func: Callable) -> Callable: Applies JIT """ + logger.debug( + "djit: func=%s, args=%s, kwargs=%s", + getattr(inner_func, "__name__", repr(inner_func)), + repr(args), + repr(kwargs), + ) + compiled = hjit( DSIG, *args, From 10da21294b5ccf8fb83344cc3e92faf8bd502a90 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 11:05:07 +0100 Subject: [PATCH 190/346] fix exports --- src/hapsira/core/jit.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 21ee45ca2..0abeda00d 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -11,6 +11,7 @@ __all__ = [ "DSIG", "hjit", + "djit", "vjit", "gjit", "sjit", From 816ddae5e0b966acaa4becae756b0de0e83d481f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 11:43:24 +0100 Subject: [PATCH 191/346] fix docs --- src/hapsira/core/events.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index 3b6d7148b..7956d056d 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -76,9 +76,9 @@ def line_of_sight_hf(r1, r2, R): Parameters ---------- - r1 : numpy.ndarray + r1 : tuple[float,float,float] The position vector of the first object with respect to a central attractor. - r2 : numpy.ndarray + r2 : tuple[float,float,float] The position vector of the second object with respect to a central attractor. R : float The radius of the central attractor. From b943b24a178ce6b100c4036a3bf00f27583e4d9f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 11:43:59 +0100 Subject: [PATCH 192/346] fix astropy depr warning --- src/hapsira/frames/ecliptic.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/hapsira/frames/ecliptic.py b/src/hapsira/frames/ecliptic.py index 65c229162..6bb4044c8 100644 --- a/src/hapsira/frames/ecliptic.py +++ b/src/hapsira/frames/ecliptic.py @@ -12,7 +12,6 @@ ) from astropy.coordinates.builtin_frames.utils import DEFAULT_OBSTIME, get_jd12 from astropy.coordinates.matrix_utilities import ( - matrix_product, matrix_transpose, rotation_matrix, ) @@ -71,7 +70,7 @@ def gcrs_to_geosolarecliptic(gcrs_coo, to_frame): rot_matrix = _make_rotation_matrix_from_reprs(sun_earth_detilt, x_axis) - return matrix_product(rot_matrix, _earth_detilt_matrix) + return rot_matrix @ _earth_detilt_matrix @frame_transform_graph.transform(DynamicMatrixTransform, GeocentricSolarEcliptic, GCRS) From 3cb926d7c2495210d5c108a0c8d6525ceeb3070d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 11:44:42 +0100 Subject: [PATCH 193/346] fix matplotlib max plots warning during tests --- tests/conftest.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index e99c92396..9835e434d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,11 +1,14 @@ from astropy import units as u from astropy.coordinates import solar_system_ephemeris from astropy.time import Time +import matplotlib as mpl import pytest from hapsira.bodies import Earth, Sun from hapsira.twobody import Orbit +mpl.rc("figure", max_open_warning=0) + solar_system_ephemeris.set("builtin") From 6320dd1a383ec152047f2029a1900e2a653b1cd8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 11:45:08 +0100 Subject: [PATCH 194/346] fix numba obj mode warn --- src/hapsira/core/math/ivp/_brentq.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index f27de4dc9..4429863a6 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -27,7 +27,7 @@ def _signbit_hf(a): return a < 0 -@jit(nopython=False) +@jit(forceobj=True) def _brentq_hf( func, # callback_type xa, # double @@ -107,7 +107,7 @@ def _brentq_hf( return xcur, funcalls, iterations, CONVERR -@jit(nopython=False) +@jit(forceobj=True) def brentq_sf( func, # func a, # double From dad958f3f388f53949cf3142480e7dc9d49b5ebe Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 11:46:22 +0100 Subject: [PATCH 195/346] do not attempt to cache dynamically generated functions, i.e. eliminate numba warnings --- src/hapsira/core/math/interpolate.py | 2 +- src/hapsira/core/thrust/change_a_inc.py | 2 +- src/hapsira/core/thrust/change_argp.py | 2 +- src/hapsira/core/thrust/change_ecc_inc.py | 2 +- .../core/thrust/change_ecc_quasioptimal.py | 2 +- src/hapsira/earth/__init__.py | 4 ++-- tests/tests_twobody/test_events.py | 2 +- tests/tests_twobody/test_perturbations.py | 18 +++++++++--------- tests/tests_twobody/test_propagation.py | 2 +- tests/tests_twobody/test_thrust.py | 12 ++++++------ 10 files changed, 24 insertions(+), 24 deletions(-) diff --git a/src/hapsira/core/math/interpolate.py b/src/hapsira/core/math/interpolate.py index a586286a9..780951ad8 100644 --- a/src/hapsira/core/math/interpolate.py +++ b/src/hapsira/core/math/interpolate.py @@ -28,7 +28,7 @@ def interp_hb(x: np.ndarray, y: np.ndarray) -> Callable: x = tuple(x) x_len = len(x) - @hjit("V(f)") + @hjit("V(f)", cache=False) def interp_hf(x_new): assert x_new >= x[0] assert x_new <= x[-1] diff --git a/src/hapsira/core/thrust/change_a_inc.py b/src/hapsira/core/thrust/change_a_inc.py index bf4f3b9b5..1027193a8 100644 --- a/src/hapsira/core/thrust/change_a_inc.py +++ b/src/hapsira/core/thrust/change_a_inc.py @@ -121,7 +121,7 @@ def change_a_inc_hb(k, a_0, a_f, inc_0, inc_f, f): k, a_0, a_f, inc_0, inc_f ) - @hjit("V(f,V,V,f)") + @hjit("V(f,V,V,f)", cache=False) def a_d_hf(t0, rr, vv, k): # Change sign of beta with the out-of-plane velocity beta_ = _beta_hf(t0, V_0, f, beta_0_) * sign_hf(rr[0] * (inc_f - inc_0)) diff --git a/src/hapsira/core/thrust/change_argp.py b/src/hapsira/core/thrust/change_argp.py index ab2d28e55..20da694a3 100644 --- a/src/hapsira/core/thrust/change_argp.py +++ b/src/hapsira/core/thrust/change_argp.py @@ -72,7 +72,7 @@ def change_argp_hb(k, a, ecc, argp_0, argp_f, f): t_f : float """ - @hjit("V(f,V,V,f)") + @hjit("V(f,V,V,f)", cache=False) def a_d_hf(t0, rr, vv, k): nu = rv2coe_hf(k, rr, vv, RV2COE_TOL)[-1] diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index 2964046b0..9a71aac91 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -134,7 +134,7 @@ def change_ecc_inc_hb(k, a, ecc_0, ecc_f, inc_0, inc_f, argp, r, v, f): ) thrust_unit = tuple(thrust_unit) - @hjit("V(f,V,V,f)") + @hjit("V(f,V,V,f)", cache=False) def a_d_hf(t0, rr, vv, k_): nu = rv2coe_hf(k_, rr, vv, RV2COE_TOL)[-1] beta_ = beta_0 * sign_hf( diff --git a/src/hapsira/core/thrust/change_ecc_quasioptimal.py b/src/hapsira/core/thrust/change_ecc_quasioptimal.py index fb1ac52d4..2fe6fe287 100644 --- a/src/hapsira/core/thrust/change_ecc_quasioptimal.py +++ b/src/hapsira/core/thrust/change_ecc_quasioptimal.py @@ -78,7 +78,7 @@ def change_ecc_quasioptimal_hb(k, a, ecc_0, ecc_f, e_vec, h_vec, r, f): ) thrust_unit = tuple(thrust_unit) - @hjit("V(f,V,V,f)") + @hjit("V(f,V,V,f)", cache=False) def a_d_hf(t0, rr, vv, k): accel_v = mul_Vs_hf(thrust_unit, f) return accel_v diff --git a/src/hapsira/earth/__init__.py b/src/hapsira/earth/__init__.py index 3cd1af6c0..09e650fc6 100644 --- a/src/hapsira/earth/__init__.py +++ b/src/hapsira/earth/__init__.py @@ -86,7 +86,7 @@ def propagate(self, tof, atmosphere=None, gravity=None, *args): J2_ = Earth.J2.value R_ = Earth.R.to_value(u.km) - @hjit("V(f,V,V,f)") + @hjit("V(f,V,V,f)", cache=False) def ad_hf(t0, rr, vv, k): return J2_perturbation_hf(t0, rr, vv, k, J2_, R_) @@ -96,7 +96,7 @@ def ad_hf(t0, rr, vv, k): def ad_hf(t0, rr, vv, k): return 0.0, 0.0, 0.0 - @djit + @djit(cache=False) def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad_vv = ad_hf(t0, rr, vv, k) diff --git a/tests/tests_twobody/test_events.py b/tests/tests_twobody/test_events.py index b43baac3b..1b0a37dc9 100644 --- a/tests/tests_twobody/test_events.py +++ b/tests/tests_twobody/test_events.py @@ -49,7 +49,7 @@ def test_altitude_crossing(): altitude_cross_event = AltitudeCrossEvent(thresh_alt, R) events = [altitude_cross_event] - @djit + @djit(cache=False) def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = atmospheric_drag_exponential_hf( diff --git a/tests/tests_twobody/test_perturbations.py b/tests/tests_twobody/test_perturbations.py index 99b7ecc62..94abf302b 100644 --- a/tests/tests_twobody/test_perturbations.py +++ b/tests/tests_twobody/test_perturbations.py @@ -41,7 +41,7 @@ def test_J2_propagation_Earth(): J2 = Earth.J2.value R_ = Earth.R.to(u.km).value - @djit + @djit(cache=False) def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = J2_perturbation_hf( @@ -132,7 +132,7 @@ def test_J3_propagation_Earth(test_params): J2 = Earth.J2.value R_ = Earth.R.to(u.km).value - @djit + @djit(cache=False) def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = J2_perturbation_hf( @@ -154,7 +154,7 @@ def f_hf(t0, rr, vv, k): J3 = Earth.J3.value - @djit + @djit(cache=False) def f_combined_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad_J2 = J2_perturbation_hf( @@ -271,7 +271,7 @@ def test_atmospheric_drag_exponential(): # dr_expected = F_r * tof (Newton's integration formula), where # F_r = -B rho(r) |r|^2 sqrt(k / |r|^3) = -B rho(r) sqrt(k |r|) - @djit + @djit(cache=False) def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = atmospheric_drag_exponential_hf( @@ -321,7 +321,7 @@ def test_atmospheric_demise(): lithobrake_event = LithobrakeEvent(R) events = [lithobrake_event] - @djit + @djit(cache=False) def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = atmospheric_drag_exponential_hf( @@ -380,7 +380,7 @@ def test_atmospheric_demise_coesa76(): lithobrake_event = LithobrakeEvent(R) events = [lithobrake_event] - @djit + @djit(cache=False) def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) @@ -617,7 +617,7 @@ def test_3rd_body_Curtis(test_params): body_r = build_ephem_interpolant(body, body_epochs) k_third = body.k.to_value(u.km**3 / u.s**2) - @djit + @djit(cache=False) def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = third_body_hf( @@ -707,7 +707,7 @@ def test_solar_pressure(t_days, deltas_expected, sun_r): ) # In Curtis, the mean distance to Sun is used. In order to validate against it, we have to do the same thing - @hjit("V(f)") + @hjit("V(f)", cache=False) def sun_normalized_hf(t0): r = sun_r(t0) # sun_r is hf, returns V return mul_Vs_hf(r, 149600000 / norm_V_hf(r)) @@ -715,7 +715,7 @@ def sun_normalized_hf(t0): R_ = Earth.R.to(u.km).value Wdivc_s = Wdivc_sun.value - @djit + @djit(cache=False) def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = radiation_pressure_hf( diff --git a/tests/tests_twobody/test_propagation.py b/tests/tests_twobody/test_propagation.py index d57d302b3..593e73db6 100644 --- a/tests/tests_twobody/test_propagation.py +++ b/tests/tests_twobody/test_propagation.py @@ -296,7 +296,7 @@ def constant_accel_hf(t0, rr, vv, k): norm_v = norm_V_hf(vv) return mul_Vs_hf(vv, accel / norm_v) - @djit + @djit(cache=False) def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = constant_accel_hf(t0, rr, vv, k) diff --git a/tests/tests_twobody/test_thrust.py b/tests/tests_twobody/test_thrust.py index 93f2300f4..cf0d1c83d 100644 --- a/tests/tests_twobody/test_thrust.py +++ b/tests/tests_twobody/test_thrust.py @@ -40,7 +40,7 @@ def test_leo_geo_numerical_safe(inc_0): s0 = Orbit.circular(Earth, a_0 - Earth.R, inc_0) # Propagate orbit - @djit + @djit(cache=False) def f_leo_geo_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = a_d_hf(t0, rr, vv, k) @@ -72,7 +72,7 @@ def test_leo_geo_numerical_fast(inc_0): s0 = Orbit.circular(Earth, a_0 * u.km - Earth.R, inc_0 * u.rad) # Propagate orbit - @djit + @djit(cache=False) def f_leo_geo_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = a_d_hf(t0, rr, vv, k) @@ -131,7 +131,7 @@ def test_sso_disposal_numerical(ecc_0, ecc_f): a_d_hf, _, t_f = change_ecc_quasioptimal(s0, ecc_f, f) # Propagate orbit - @djit + @djit(cache=False) def f_ss0_disposal_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = a_d_hf(t0, rr, vv, k) @@ -202,7 +202,7 @@ def test_geo_cases_numerical(ecc_0, ecc_f): a_d_hf, _, t_f = change_ecc_inc(orb_0=s0, ecc_f=ecc_f, inc_f=inc_f, f=f) # Propagate orbit - @djit + @djit(cache=False) def f_geo_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = a_d_hf(t0, rr, vv, k) @@ -285,7 +285,7 @@ def test_soyuz_standard_gto_numerical_safe(): ) # Propagate orbit - @djit + @djit(cache=False) def f_soyuz_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = a_d_hf(t0, rr, vv, k) @@ -323,7 +323,7 @@ def test_soyuz_standard_gto_numerical_fast(): ) # Propagate orbit - @djit + @djit(cache=False) def f_soyuz_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) du_ad = a_d_hf(t0, rr, vv, k) From b0bc653872f21d03901aa035ec5249d3a9bd3a98 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 11:58:42 +0100 Subject: [PATCH 196/346] pylint --- src/hapsira/twobody/events.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index 2755c4df0..a47af82cd 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -281,5 +281,9 @@ def __call__(self, t, u_, k): pos_coord = self._pos_coords.pop(0) if self._pos_coords else self._last_coord # Need to cast `pos_coord` to array since `norm` inside numba only works for arrays, not lists. - delta_angle = line_of_sight_gf(u_[:3], np.array(pos_coord), self._R) + delta_angle = line_of_sight_gf( # pylint: disable=E1120 + u_[:3], + np.array(pos_coord), + self._R, + ) return delta_angle From c818264ee006c3726a70db5e7ba6d3db0696829b Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 12:45:36 +0100 Subject: [PATCH 197/346] less wrapping for events --- src/hapsira/core/math/ivp/_solve.py | 44 ++++++++++++++++++++++------- 1 file changed, 34 insertions(+), 10 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 063ab281b..89cf8cee0 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -13,7 +13,9 @@ ] -def _solve_event_equation(event, sol, t_old, t): +def _solve_event_equation( + event: Callable, sol: Callable, t_old: float, t: float, argk: float +) -> float: """Solve an equation corresponding to an ODE event. The equation is ``event(t, y(t)) = 0``, here ``y(t)`` is known from an @@ -37,10 +39,27 @@ def _solve_event_equation(event, sol, t_old, t): Found solution. """ - return brentq(lambda t: event(t, sol(t)), t_old, t, xtol=4 * EPS, rtol=4 * EPS) - - -def _handle_events(sol, events, active_events, is_terminal, t_old, t): + def wrapper(t): + return event(t, sol(t), argk) + + return brentq( + wrapper, + t_old, + t, + xtol=4 * EPS, + rtol=4 * EPS, + ) + + +def _handle_events( + sol, + events: List[Callable], + active_events, + is_terminal, + t_old: float, + t: float, + argk: float, +): """Helper function to handle events. Parameters @@ -68,7 +87,7 @@ def _handle_events(sol, events, active_events, is_terminal, t_old, t): Whether a terminal event occurred. """ roots = [ - _solve_event_equation(events[event_index], sol, t_old, t) + _solve_event_equation(events[event_index], sol, t_old, t, argk) for event_index in active_events ] @@ -300,8 +319,7 @@ def solve_ivp( events, is_terminal, event_dir = _prepare_events(events) if events is not None: - events = [lambda t, x, event=event: event(t, x, argk) for event in events] - g = [event(t0, y0) for event in events] + g = [event(t0, y0, argk) for event in events] status = None while status is None: @@ -321,11 +339,17 @@ def solve_ivp( interpolants.append(sol) if events is not None: - g_new = [event(t, y) for event in events] + g_new = [event(t, y, argk) for event in events] active_events = _find_active_events(g, g_new, event_dir) if active_events.size > 0: _, roots, terminate = _handle_events( - sol, events, active_events, is_terminal, t_old, t + sol, + events, + active_events, + is_terminal, + t_old, + t, + argk, ) if terminate: status = 1 From 51f7eeb658160c5d99be88f43359c52029965d89 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 13:22:10 +0100 Subject: [PATCH 198/346] more r and v in solver --- src/hapsira/core/math/ivp/_rk.py | 86 ++++++++++++++++---------- src/hapsira/core/math/ivp/_solve.py | 8 ++- src/hapsira/core/propagation/cowell.py | 11 +--- 3 files changed, 61 insertions(+), 44 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 3c235be2e..ce693705d 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -252,15 +252,14 @@ def __init__( self, fun: Callable, t0: float, - y0: np.array, + rr: tuple, + vv: tuple, t_bound: float, argk: float, max_step: float = np.inf, rtol: float = 1e-3, atol: float = 1e-6, ): - assert y0.shape == (N_RV,) - assert np.isfinite(y0).all() assert max_step > 0 assert atol >= 0 @@ -268,7 +267,8 @@ def __init__( rtol = 100 * EPS self.t = t0 - self.y = y0 + self.rr = rr + self.vv = vv self.t_bound = t_bound self.max_step = max_step self.fun = fun @@ -278,30 +278,32 @@ def __init__( self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 - self.K_extended = np.empty((N_STAGES_EXTENDED, N_RV), dtype=self.y.dtype) + self.K_extended = np.empty( + (N_STAGES_EXTENDED, N_RV), dtype=float + ) # TODO set type self.K = self.K_extended[: N_STAGES + 1, :] - self.y_old = None + self.rr_old = None + self.vv_old = None self.t_old = None self.h_previous = None self.status = "running" - rr, vv = self.fun( + self.fr, self.fv = self.fun( self.t, - array_to_V_hf(self.y[:3]), - array_to_V_hf(self.y[3:]), + self.rr, + self.vv, self.argk, ) # TODO call into hf - self.f = np.array([*rr, *vv]) self.h_abs = _select_initial_step_hf( self.fun, self.t, - array_to_V_hf(self.y[:3]), - array_to_V_hf(self.y[3:]), + self.rr, + self.vv, self.argk, - array_to_V_hf(self.f[:3]), - array_to_V_hf(self.f[3:]), + self.fr, + self.fv, self.direction, ERROR_ESTIMATOR_ORDER, self.rtol, @@ -333,10 +335,10 @@ def step(self): self.fun, self.argk, self.t, - array_to_V_hf(self.y[:3]), - array_to_V_hf(self.y[3:]), - array_to_V_hf(self.f[:3]), - array_to_V_hf(self.f[3:]), + self.rr, + self.vv, + self.fr, + self.fv, self.max_step, self.rtol, self.atol, @@ -349,11 +351,12 @@ def step(self): if success: self.h_previous = rets[0] # self.y_old = np.array([*rets[1], *rets[2]]) - self.y_old = self.y + self.rr_old = self.rr + self.vv_old = self.vv self.t = rets[1] - self.y = np.array([*rets[2], *rets[3]]) + self.rr, self.vv = rets[2], rets[3] self.h_abs = rets[4] - self.f = np.array([*rets[5], *rets[6]]) + self.fr, self.fv = rets[5], rets[6] self.K[: N_STAGES + 1, :N_RV] = np.array(rets[7]) if not success: @@ -374,31 +377,48 @@ def dense_output(self): sol : `DenseOutput` Local interpolant over the last successful step. """ - assert self.t_old is not None + assert self.t_old is not None assert self.t != self.t_old K = self.K_extended h = self.h_previous + for s, (a, c) in enumerate(zip(self.A_EXTRA, self.C_EXTRA), start=N_STAGES + 1): dy = np.dot(K[:s].T, a[:s]) * h - y_ = self.y_old + dy + rr_ = add_VV_hf(self.rr_old, array_to_V_hf(dy[:3])) + vv_ = add_VV_hf(self.vv_old, array_to_V_hf(dy[3:])) rr, vv = self.fun( self.t_old + c * h, - array_to_V_hf(y_[:3]), - array_to_V_hf(y_[3:]), + rr_, + vv_, self.argk, ) # TODO call into hf K[s] = np.array([*rr, *vv]) - F = np.empty((INTERPOLATOR_POWER, N_RV), dtype=self.y_old.dtype) + F = np.empty((INTERPOLATOR_POWER, N_RV), dtype=float) # TODO use correct type + + fr_old = array_to_V_hf(K[0, :3]) + fv_old = array_to_V_hf(K[0, 3:]) + + delta_rr = sub_VV_hf(self.rr, self.rr_old) + delta_vv = sub_VV_hf(self.vv, self.vv_old) - f_old = K[0] - delta_y = self.y - self.y_old + F[0, :3] = delta_rr + F[0, 3:] = delta_vv - F[0] = delta_y - F[1] = h * f_old - delta_y - F[2] = 2 * delta_y - h * (self.f + f_old) - F[3:] = h * np.dot(self.D, K) + F[1, :3] = sub_VV_hf(mul_Vs_hf(fr_old, h), delta_rr) + F[1, 3:] = sub_VV_hf(mul_Vs_hf(fv_old, h), delta_vv) - return Dop853DenseOutput(self.t_old, self.t, self.y_old, F) + F[2, :3] = sub_VV_hf( + mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(self.fr, fr_old), h) + ) + F[2, 3:] = sub_VV_hf( + mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(self.fv, fv_old), h) + ) + + F[3:, :] = h * np.dot(self.D, K) # TODO + + return Dop853DenseOutput( + self.t_old, self.t, np.array([*self.rr_old, *self.vv_old]), F + ) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 89cf8cee0..7488287d0 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -163,7 +163,8 @@ def solve_ivp( fun: Callable, t0: float, tf: float, - y0: np.ndarray, # (6,) + rr: Tuple[float, float, float], + vv: Tuple[float, float, float], argk: float, events: Optional[List[Callable]] = None, **options, @@ -310,8 +311,9 @@ def solve_ivp( """ - solver = DOP853(fun, t0, y0, tf, argk, **options) + solver = DOP853(fun, t0, rr, vv, tf, argk, **options) + y0 = np.array([*rr, *vv]) # TODO turn into tuples ts = [t0] interpolants = [] @@ -333,7 +335,7 @@ def solve_ivp( t_old = solver.t_old t = solver.t - y = solver.y + y = np.array([*solver.rr, *solver.vv]) # TODO turn into tuples sol = solver.dense_output() interpolants.append(sol) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index be5dd45e9..5766fb7a6 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -1,5 +1,4 @@ -import numpy as np - +from ..jit import array_to_V_hf from ..math.ivp import solve_ivp from ..propagation.base import func_twobody_hf @@ -26,16 +25,12 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=None, f=func_twobody_hf assert hasattr(f, "djit") # DEBUG check for compiler flag assert isinstance(rtol, float) - x, y, z = r - vx, vy, vz = v - - u0 = np.array([x, y, z, vx, vy, vz]) - sol, success = solve_ivp( f, 0.0, float(max(tofs)), - u0, + array_to_V_hf(r), + array_to_V_hf(v), argk=k, rtol=rtol, atol=atol, From 980f84d749ecffb418758682321e339fc07a27e2 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 13:40:32 +0100 Subject: [PATCH 199/346] cleanup args --- src/hapsira/core/math/ivp/_rk.py | 14 +--- src/hapsira/core/math/ivp/_solve.py | 121 +--------------------------- 2 files changed, 8 insertions(+), 127 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index ce693705d..73d5f7ac4 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -90,15 +90,13 @@ def _select_initial_step_hf( @hjit( f"Tuple([b1,f,f,V,V,f,V,V,{KSIG:s}])" - f"(F({DSIG:s}),f,f,V,V,V,V,f,f,f,f,f,f,{KSIG:s})" + f"(F({DSIG:s}),f,f,V,V,V,V,f,f,f,f,f,{KSIG:s})" ) def _step_impl_hf( - fun, argk, t, rr, vv, fr, fv, max_step, rtol, atol, direction, h_abs, t_bound, K + fun, argk, t, rr, vv, fr, fv, rtol, atol, direction, h_abs, t_bound, K ): min_step = 10 * abs(nextafter_hf(t, direction * inf) - t) - if h_abs > max_step: - h_abs = max_step if h_abs < min_step: h_abs = min_step @@ -256,11 +254,9 @@ def __init__( vv: tuple, t_bound: float, argk: float, - max_step: float = np.inf, - rtol: float = 1e-3, - atol: float = 1e-6, + rtol: float, + atol: float, ): - assert max_step > 0 assert atol >= 0 if rtol < 100 * EPS: @@ -270,7 +266,6 @@ def __init__( self.rr = rr self.vv = vv self.t_bound = t_bound - self.max_step = max_step self.fun = fun self.argk = argk self.rtol = rtol @@ -339,7 +334,6 @@ def step(self): self.vv, self.fr, self.fv, - self.max_step, self.rtol, self.atol, self.direction, diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 7488287d0..ec1708b66 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -166,8 +166,9 @@ def solve_ivp( rr: Tuple[float, float, float], vv: Tuple[float, float, float], argk: float, + rtol: float, + atol: float, events: Optional[List[Callable]] = None, - **options, ) -> Tuple[OdeSolution, bool]: """Solve an initial value problem for a system of ODEs. @@ -186,15 +187,6 @@ def solve_ivp( y0 : array_like, shape (n,) Initial state. For problems in the complex domain, pass `y0` with a complex data type (even if the initial value is purely real). - method : string or `OdeSolver`, optional - Integration method to use: - * 'DOP853': Explicit Runge-Kutta method of order 8 [13]_. - Python implementation of the "DOP853" algorithm originally - written in Fortran [14]_. A 7-th order interpolation polynomial - accurate to 7-th order is used for the dense output. - Can be applied in the complex domain. - You can also pass an arbitrary class derived from `OdeSolver` which - implements the solver. events : callable, or list of callables, optional Events to track. If None (default), no events will be tracked. Each event occurs at the zeros of a continuous function of time and @@ -203,115 +195,10 @@ def solve_ivp( ``event(t, y(t)) = 0`` using a root-finding algorithm. By default, all zeros will be found. The solver looks for a sign change over each step, so if multiple zero crossings occur within one step, events may be - missed. Additionally each `event` function might have the following - attributes: - - terminal: bool, optional - Whether to terminate integration if this event occurs. - Implicitly False if not assigned. - direction: float, optional - Direction of a zero crossing. If `direction` is positive, - `event` will only trigger when going from negative to positive, - and vice versa if `direction` is negative. If 0, then either - direction will trigger event. Implicitly 0 if not assigned. - - You can assign attributes like ``event.terminal = True`` to any - function in Python. - **options - Options passed to a chosen solver. All options available for already - implemented solvers are listed below. - first_step : float or None, optional - Initial step size. Default is `None` which means that the algorithm - should choose. - max_step : float, optional - Maximum allowed step size. Default is np.inf, i.e., the step size is not - bounded and determined solely by the solver. - rtol, atol : float or array_like, optional - Relative and absolute tolerances. The solver keeps the local error - estimates less than ``atol + rtol * abs(y)``. Here `rtol` controls a - relative accuracy (number of correct digits), while `atol` controls - absolute accuracy (number of correct decimal places). To achieve the - desired `rtol`, set `atol` to be smaller than the smallest value that - can be expected from ``rtol * abs(y)`` so that `rtol` dominates the - allowable error. If `atol` is larger than ``rtol * abs(y)`` the - number of correct digits is not guaranteed. Conversely, to achieve the - desired `atol` set `rtol` such that ``rtol * abs(y)`` is always smaller - than `atol`. If components of y have different scales, it might be - beneficial to set different `atol` values for different components by - passing array_like with shape (n,) for `atol`. Default values are - 1e-3 for `rtol` and 1e-6 for `atol`. - jac : array_like, sparse_matrix, callable or None, optional - Jacobian matrix of the right-hand side of the system with respect - to y, required by the 'Radau', 'BDF' and 'LSODA' method. The - Jacobian matrix has shape (n, n) and its element (i, j) is equal to - ``d f_i / d y_j``. There are three ways to define the Jacobian: - - * If array_like or sparse_matrix, the Jacobian is assumed to - be constant. Not supported by 'LSODA'. - * If callable, the Jacobian is assumed to depend on both - t and y; it will be called as ``jac(t, y)``, as necessary. - For 'Radau' and 'BDF' methods, the return value might be a - sparse matrix. - * If None (default), the Jacobian will be approximated by - finite differences. - - It is generally recommended to provide the Jacobian rather than - relying on a finite-difference approximation. - jac_sparsity : array_like, sparse matrix or None, optional - Defines a sparsity structure of the Jacobian matrix for a finite- - difference approximation. Its shape must be (n, n). This argument - is ignored if `jac` is not `None`. If the Jacobian has only few - non-zero elements in *each* row, providing the sparsity structure - will greatly speed up the computations [10]_. A zero entry means that - a corresponding element in the Jacobian is always zero. If None - (default), the Jacobian is assumed to be dense. - Not supported by 'LSODA', see `lband` and `uband` instead. - lband, uband : int or None, optional - Parameters defining the bandwidth of the Jacobian for the 'LSODA' - method, i.e., ``jac[i, j] != 0 only for i - lband <= j <= i + uband``. - Default is None. Setting these requires your jac routine to return the - Jacobian in the packed format: the returned array must have ``n`` - columns and ``uband + lband + 1`` rows in which Jacobian diagonals are - written. Specifically ``jac_packed[uband + i - j , j] = jac[i, j]``. - The same format is used in `scipy.linalg.solve_banded` (check for an - illustration). These parameters can be also used with ``jac=None`` to - reduce the number of Jacobian elements estimated by finite differences. - min_step : float, optional - The minimum allowed step size for 'LSODA' method. - By default `min_step` is zero. - - Returns - ------- - Bunch object with the following fields defined: - t : ndarray, shape (n_points,) - Time points. - y : ndarray, shape (n, n_points) - Values of the solution at `t`. - sol : `OdeSolution` or None - Found solution as `OdeSolution` instance; None if `dense_output` was - set to False. - t_events : list of ndarray or None - Contains for each event type a list of arrays at which an event of - that type event was detected. None if `events` was None. - y_events : list of ndarray or None - For each value of `t_events`, the corresponding value of the solution. - None if `events` was None. - nlu : int - Number of LU decompositions. - status : int - Reason for algorithm termination: - - * -1: Integration step failed. - * 0: The solver successfully reached the end of `tspan`. - * 1: A termination event occurred. - - success : bool - True if the solver reached the interval end or a termination event - occurred (``status >= 0``). - + missed. """ - solver = DOP853(fun, t0, rr, vv, tf, argk, **options) + solver = DOP853(fun, t0, rr, vv, tf, argk, rtol, atol) y0 = np.array([*rr, *vv]) # TODO turn into tuples ts = [t0] From bda6ba9e740dc29d944bea870013d3def53221e6 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 14:06:21 +0100 Subject: [PATCH 200/346] compile threebody helper --- src/hapsira/threebody/restricted.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/hapsira/threebody/restricted.py b/src/hapsira/threebody/restricted.py index 67ffdf66d..772871d06 100644 --- a/src/hapsira/threebody/restricted.py +++ b/src/hapsira/threebody/restricted.py @@ -7,6 +7,7 @@ from astropy import units as u import numpy as np +from hapsira.core.jit import hjit from hapsira.core.math.ivp import brentq from hapsira.util import norm @@ -37,6 +38,7 @@ def lagrange_points(r12, m1, m2): """ pi2 = (m2 / (m1 + m2)).value + @hjit("f(f)", cache=False) def eq_L123(xi): aux = (1 - pi2) * (xi + pi2) / abs(xi + pi2) ** 3 aux += pi2 * (xi + pi2 - 1) / abs(xi + pi2 - 1) ** 3 From f6117c83f84c1b14f153eea213e2bd90cbdab9bc Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 14:17:46 +0100 Subject: [PATCH 201/346] cleanup --- src/hapsira/core/math/ivp/_brentq.py | 59 +++++++--------------------- 1 file changed, 14 insertions(+), 45 deletions(-) diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index 4429863a6..48cd32cbd 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -1,4 +1,4 @@ -from math import fabs +from math import fabs, isnan import operator @@ -41,21 +41,18 @@ def _brentq_hf( fpre, fcur, fblk = 0.0, 0.0, 0.0 spre, scur = 0.0, 0.0 - iterations = 0 - fpre = func(xpre) + assert not isnan(fpre) fcur = func(xcur) - funcalls = 2 + assert not isnan(fcur) if fpre == 0: - return xpre, funcalls, iterations, CONVERGED + return xpre, CONVERGED if fcur == 0: - return xcur, funcalls, iterations, CONVERGED + return xcur, CONVERGED if _signbit_hf(fpre) == _signbit_hf(fcur): - return 0.0, funcalls, iterations, SIGNERR + return 0.0, SIGNERR - iterations = 0 for _ in range(0, iter_): - iterations += 1 if fpre != 0 and fcur != 0 and _signbit_hf(fpre) != _signbit_hf(fcur): xblk = xpre fblk = fpre @@ -73,7 +70,7 @@ def _brentq_hf( delta = (xtol + rtol * fabs(xcur)) / 2 sbis = (xblk - xcur) / 2 if fcur == 0 or fabs(sbis) < delta: - return xcur, funcalls, iterations, CONVERGED + return xcur, CONVERGED if fabs(spre) > delta and fabs(fcur) < fabs(fpre): if xpre == xblk: @@ -102,9 +99,9 @@ def _brentq_hf( xcur += delta if sbis > 0 else -delta fcur = func(xcur) - funcalls += 1 + assert not isnan(fcur) - return xcur, funcalls, iterations, CONVERR + return xcur, CONVERR @jit(forceobj=True) @@ -116,12 +113,7 @@ def brentq_sf( rtol, # double iter_, # int ): - if xtol < 0: - raise ValueError("xtol must be >= 0") - if iter_ < 0: - raise ValueError("maxiter should be > 0") - - zero, funcalls, iterations, error_num = _brentq_hf( + zero, error_num = _brentq_hf( func, a, b, @@ -129,31 +121,10 @@ def brentq_sf( rtol, iter_, ) - - if error_num == SIGNERR: - raise ValueError("f(a) and f(b) must have different signs") - if error_num == CONVERR: - raise RuntimeError("Failed to converge after %d iterations." % iterations) - + assert error_num == CONVERGED return zero # double -def _wrap_nan_raise(f): - def f_raise(x): - fx = f(x) - f_raise._function_calls += 1 - if np.isnan(fx): - msg = f"The function value at x={x} is NaN; " "solver cannot continue." - err = ValueError(msg) - err._x = x - err._function_calls = f_raise._function_calls - raise err - return fx - - f_raise._function_calls = 0 - return f_raise - - def brentq( f, a, @@ -167,10 +138,8 @@ def brentq( https://github.com/scipy/scipy/blob/d23363809572e9a44074a3f06f66137083446b48/scipy/optimize/_zeros_py.py#L682 """ maxiter = operator.index(maxiter) - if xtol <= 0: - raise ValueError("xtol too small (%g <= 0)" % xtol) - if rtol < BRENTQ_RTOL: - raise ValueError(f"rtol too small ({rtol:g} < {BRENTQ_RTOL:g})") - f = _wrap_nan_raise(f) + assert xtol > 0 + assert rtol >= BRENTQ_RTOL + assert maxiter >= 0 r = brentq_sf(f, a, b, xtol, rtol, maxiter) return r From f6e7adcdd51089b4d5527ace6fb35bd3b84d844e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 14:30:07 +0100 Subject: [PATCH 202/346] cleanup --- src/hapsira/core/math/ivp/_brentq.py | 65 +++++++++++----------------- 1 file changed, 26 insertions(+), 39 deletions(-) diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index 48cd32cbd..626622aea 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -1,33 +1,29 @@ from math import fabs, isnan -import operator - -from numba import njit, jit -import numpy as np +from ..linalg import EPS +from ...jit import hjit CONVERGED = 0 SIGNERR = -1 CONVERR = -2 -# EVALUEERR = -3 -INPROGRESS = 1 BRENTQ_ITER = 100 BRENTQ_XTOL = 2e-12 -BRENTQ_RTOL = 4 * np.finfo(float).eps +BRENTQ_RTOL = 4 * EPS -@njit -def _min_hf(a, b): +@hjit("f(f,f)") +def _min_ss_hf(a, b): return a if a < b else b -@njit -def _signbit_hf(a): +@hjit("b1(f)") +def _signbit_s_hf(a): return a < 0 -@jit(forceobj=True) +@hjit("f(F(f(f)),f,f,f,f,f)", forceobj=True, nopython=False, cache=False) def _brentq_hf( func, # callback_type xa, # double @@ -49,11 +45,11 @@ def _brentq_hf( return xpre, CONVERGED if fcur == 0: return xcur, CONVERGED - if _signbit_hf(fpre) == _signbit_hf(fcur): + if _signbit_s_hf(fpre) == _signbit_s_hf(fcur): return 0.0, SIGNERR for _ in range(0, iter_): - if fpre != 0 and fcur != 0 and _signbit_hf(fpre) != _signbit_hf(fcur): + if fpre != 0 and fcur != 0 and _signbit_s_hf(fpre) != _signbit_s_hf(fcur): xblk = xpre fblk = fpre scur = xcur - xpre @@ -81,7 +77,7 @@ def _brentq_hf( stry = ( -fcur * (fblk * dblk - fpre * dpre) / (dblk * dpre * (fblk - fpre)) ) - if 2 * fabs(stry) < _min_hf(fabs(spre), 3 * fabs(sbis) - delta): + if 2 * fabs(stry) < _min_ss_hf(fabs(spre), 3 * fabs(sbis) - delta): spre = scur scur = stry else: @@ -104,27 +100,7 @@ def _brentq_hf( return xcur, CONVERR -@jit(forceobj=True) -def brentq_sf( - func, # func - a, # double - b, # double - xtol, # double - rtol, # double - iter_, # int -): - zero, error_num = _brentq_hf( - func, - a, - b, - xtol, - rtol, - iter_, - ) - assert error_num == CONVERGED - return zero # double - - +@hjit("f(F(f(f)),f,f,f,f,f)", forceobj=True, nopython=False, cache=False) def brentq( f, a, @@ -137,9 +113,20 @@ def brentq( Loosely adapted from https://github.com/scipy/scipy/blob/d23363809572e9a44074a3f06f66137083446b48/scipy/optimize/_zeros_py.py#L682 """ - maxiter = operator.index(maxiter) + assert xtol > 0 assert rtol >= BRENTQ_RTOL assert maxiter >= 0 - r = brentq_sf(f, a, b, xtol, rtol, maxiter) - return r + + zero, error_num = _brentq_hf( + f, + a, + b, + xtol, + rtol, + maxiter, + ) + + assert error_num == CONVERGED + + return zero From d4b1ab3f8321e5718cafcf3e2176058bddd4034c Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 16:09:08 +0100 Subject: [PATCH 203/346] cleanup --- src/hapsira/core/math/ivp/__init__.py | 20 +++++- src/hapsira/core/math/ivp/_brentq.py | 94 +++++++++++++-------------- src/hapsira/core/math/ivp/_solve.py | 11 ++-- src/hapsira/threebody/restricted.py | 23 +++++-- 4 files changed, 90 insertions(+), 58 deletions(-) diff --git a/src/hapsira/core/math/ivp/__init__.py b/src/hapsira/core/math/ivp/__init__.py index 2aec5b6bf..138cec664 100644 --- a/src/hapsira/core/math/ivp/__init__.py +++ b/src/hapsira/core/math/ivp/__init__.py @@ -1,7 +1,23 @@ from ._solve import solve_ivp -from ._brentq import brentq +from ._brentq import ( + BRENTQ_CONVERGED, + BRENTQ_SIGNERR, + BRENTQ_CONVERR, + BRENTQ_ERROR, + BRENTQ_XTOL, + BRENTQ_RTOL, + BRENTQ_MAXITER, + brentq_hf, +) __all__ = [ "solve_ivp", - "brentq", + "BRENTQ_CONVERGED", + "BRENTQ_SIGNERR", + "BRENTQ_CONVERR", + "BRENTQ_ERROR", + "BRENTQ_XTOL", + "BRENTQ_RTOL", + "BRENTQ_MAXITER", + "brentq_hf", ] diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index 626622aea..af85d891a 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -4,13 +4,26 @@ from ...jit import hjit -CONVERGED = 0 -SIGNERR = -1 -CONVERR = -2 +__all__ = [ + "BRENTQ_CONVERGED", + "BRENTQ_SIGNERR", + "BRENTQ_CONVERR", + "BRENTQ_ERROR", + "BRENTQ_XTOL", + "BRENTQ_RTOL", + "BRENTQ_MAXITER", + "brentq_hf", +] + + +BRENTQ_CONVERGED = 0 +BRENTQ_SIGNERR = -1 +BRENTQ_CONVERR = -2 +BRENTQ_ERROR = -3 -BRENTQ_ITER = 100 BRENTQ_XTOL = 2e-12 BRENTQ_RTOL = 4 * EPS +BRENTQ_MAXITER = 100 @hjit("f(f,f)") @@ -23,32 +36,48 @@ def _signbit_s_hf(a): return a < 0 -@hjit("f(F(f(f)),f,f,f,f,f)", forceobj=True, nopython=False, cache=False) -def _brentq_hf( +@hjit("Tuple([f,i8])(F(f(f)),f,f,f,f,f)", forceobj=True, nopython=False, cache=False) +def brentq_hf( func, # callback_type xa, # double xb, # double xtol, # double rtol, # double - iter_, # int + maxiter, # int ): + """ + Loosely adapted from + https://github.com/scipy/scipy/blob/d23363809572e9a44074a3f06f66137083446b48/scipy/optimize/_zeros_py.py#L682 + """ + + # if not xtol + 0. > 0: + # return 0., BRENTQ_ERROR + # if not rtol + 0. >= BRENTQ_RTOL: + # return 0., BRENTQ_ERROR + # if not maxiter + 0 >= 0: + # return 0., BRENTQ_ERROR + xpre, xcur = xa, xb xblk = 0.0 fpre, fcur, fblk = 0.0, 0.0, 0.0 spre, scur = 0.0, 0.0 fpre = func(xpre) - assert not isnan(fpre) + if isnan(fpre): + return 0.0, BRENTQ_ERROR + fcur = func(xcur) - assert not isnan(fcur) + if isnan(fcur): + return 0.0, BRENTQ_ERROR + if fpre == 0: - return xpre, CONVERGED + return xpre, BRENTQ_CONVERGED if fcur == 0: - return xcur, CONVERGED + return xcur, BRENTQ_CONVERGED if _signbit_s_hf(fpre) == _signbit_s_hf(fcur): - return 0.0, SIGNERR + return 0.0, BRENTQ_SIGNERR - for _ in range(0, iter_): + for _ in range(0, maxiter): if fpre != 0 and fcur != 0 and _signbit_s_hf(fpre) != _signbit_s_hf(fcur): xblk = xpre fblk = fpre @@ -66,7 +95,7 @@ def _brentq_hf( delta = (xtol + rtol * fabs(xcur)) / 2 sbis = (xblk - xcur) / 2 if fcur == 0 or fabs(sbis) < delta: - return xcur, CONVERGED + return xcur, BRENTQ_CONVERGED if fabs(spre) > delta and fabs(fcur) < fabs(fpre): if xpre == xblk: @@ -95,38 +124,7 @@ def _brentq_hf( xcur += delta if sbis > 0 else -delta fcur = func(xcur) - assert not isnan(fcur) - - return xcur, CONVERR - - -@hjit("f(F(f(f)),f,f,f,f,f)", forceobj=True, nopython=False, cache=False) -def brentq( - f, - a, - b, - xtol=BRENTQ_XTOL, - rtol=BRENTQ_RTOL, - maxiter=BRENTQ_ITER, -): - """ - Loosely adapted from - https://github.com/scipy/scipy/blob/d23363809572e9a44074a3f06f66137083446b48/scipy/optimize/_zeros_py.py#L682 - """ - - assert xtol > 0 - assert rtol >= BRENTQ_RTOL - assert maxiter >= 0 - - zero, error_num = _brentq_hf( - f, - a, - b, - xtol, - rtol, - maxiter, - ) - - assert error_num == CONVERGED + if isnan(fcur): + return 0.0, BRENTQ_ERROR - return zero + return xcur, BRENTQ_CONVERR diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index ec1708b66..b35e8aced 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -2,7 +2,7 @@ import numpy as np -from ._brentq import brentq +from ._brentq import brentq_hf, BRENTQ_CONVERGED, BRENTQ_MAXITER from ._solution import OdeSolution from ._rk import DOP853 from ...math.linalg import EPS @@ -42,13 +42,16 @@ def _solve_event_equation( def wrapper(t): return event(t, sol(t), argk) - return brentq( + value, status = brentq_hf( wrapper, t_old, t, - xtol=4 * EPS, - rtol=4 * EPS, + 4 * EPS, + 4 * EPS, + BRENTQ_MAXITER, ) + assert BRENTQ_CONVERGED == status + return value def _handle_events( diff --git a/src/hapsira/threebody/restricted.py b/src/hapsira/threebody/restricted.py index 772871d06..1addf91e2 100644 --- a/src/hapsira/threebody/restricted.py +++ b/src/hapsira/threebody/restricted.py @@ -8,7 +8,13 @@ import numpy as np from hapsira.core.jit import hjit -from hapsira.core.math.ivp import brentq +from hapsira.core.math.ivp import ( + brentq_hf, + BRENTQ_XTOL, + BRENTQ_RTOL, + BRENTQ_MAXITER, + BRENTQ_CONVERGED, +) from hapsira.util import norm @@ -51,15 +57,24 @@ def eq_L123(xi): tol = 1e-11 # `brentq` uses a xtol of 2e-12, so it should be covered a = -pi2 + tol b = 1 - pi2 - tol - xi = brentq(eq_L123, a, b) + xi, status = brentq_hf( + eq_L123, a, b, BRENTQ_XTOL, BRENTQ_RTOL, BRENTQ_MAXITER + ) # TODO call into hf + assert status == BRENTQ_CONVERGED lp[0] = xi + pi2 # L2 - xi = brentq(eq_L123, 1, 1.5) + xi, status = brentq_hf( + eq_L123, 1, 1.5, BRENTQ_XTOL, BRENTQ_RTOL, BRENTQ_MAXITER + ) # TODO call into hf + assert status == BRENTQ_CONVERGED lp[1] = xi + pi2 # L3 - xi = brentq(eq_L123, -1.5, -1) + xi, status = brentq_hf( + eq_L123, -1.5, -1, BRENTQ_XTOL, BRENTQ_RTOL, BRENTQ_MAXITER + ) # TODO call into hf + assert status == BRENTQ_CONVERGED lp[2] = xi + pi2 # L4, L5 From 47f49e8b0221201b5fd07605034432d4637b5400 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 18:58:49 +0100 Subject: [PATCH 204/346] events cleanup --- src/hapsira/twobody/events.py | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index a47af82cd..f4bfd7a57 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -1,3 +1,4 @@ +from abc import ABC, abstractmethod from warnings import warn from astropy import units as u @@ -14,7 +15,20 @@ ) -class Event: +__all__ = [ + "BaseEvent", + "AltitudeCrossEvent", + "LithobrakeEvent", + "LatitudeCrossEvent", + "EclipseEvent", + "PenumbraEvent", + "UmbraEvent", + "NodeCrossEvent", + "LosEvent", +] + + +class BaseEvent(ABC): """Base class for event functionalities. Parameters @@ -42,11 +56,12 @@ def direction(self): def last_t(self): return self._last_t << u.s + @abstractmethod def __call__(self, t, u, k): raise NotImplementedError -class AltitudeCrossEvent(Event): +class AltitudeCrossEvent(BaseEvent): """Detect if a satellite crosses a specific threshold altitude. Parameters @@ -94,7 +109,7 @@ def __init__(self, R, terminal=True): super().__init__(0, R, terminal, direction=-1) -class LatitudeCrossEvent(Event): +class LatitudeCrossEvent(BaseEvent): """Detect if a satellite crosses a specific threshold latitude. Parameters @@ -127,7 +142,7 @@ def __call__(self, t, u_, k): return np.rad2deg(lat_) - self._lat -class EclipseEvent(Event): +class EclipseEvent(BaseEvent): """Base class for the eclipse event. Parameters @@ -225,7 +240,7 @@ def __call__(self, t, u_, k): return shadow_function -class NodeCrossEvent(Event): +class NodeCrossEvent(BaseEvent): """Detect equatorial node (ascending or descending) crossings. Parameters @@ -248,7 +263,7 @@ def __call__(self, t, u_, k): return u_[2] -class LosEvent(Event): +class LosEvent(BaseEvent): """Detect whether there exists a LOS between two satellites. Parameters From 0331c58602ffdc27b6a3f636f027b3afb337765e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 20:29:24 +0100 Subject: [PATCH 205/346] jit eclipse_function --- src/hapsira/core/events.py | 43 ++++++++++++++++++----------------- src/hapsira/twobody/events.py | 24 ++++++++++++------- 2 files changed, 38 insertions(+), 29 deletions(-) diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index 7956d056d..336a21f98 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -1,33 +1,35 @@ -from math import acos +from math import acos, cos, sin from numba import njit as jit import numpy as np from .elements import coe_rotation_matrix_hf, rv2coe_hf, RV2COE_TOL from .jit import array_to_V_hf, hjit, gjit -from .math.linalg import norm_V_hf, matmul_VV_hf +from .math.linalg import matmul_VV_hf, norm_V_hf from .util import planetocentric_to_AltAz_hf __all__ = [ - "eclipse_function", + "eclipse_function_hf", "line_of_sight_hf", "line_of_sight_gf", "elevation_function", ] -@jit -def eclipse_function(k, u_, r_sec, R_sec, R_primary, umbra=True): +@hjit("f(f,V,V,V,f,f,b1)") +def eclipse_function_hf(k, rr, vv, r_sec, R_sec, R_primary, umbra): """Calculates a continuous shadow function. Parameters ---------- k : float Standard gravitational parameter (km^3 / s^2). - u_ : numpy.ndarray - Satellite position and velocity vector with respect to the primary body. - r_sec : numpy.ndarray + rr : tuple[float,float,float] + Satellite position vector with respect to the primary body. + vv : tuple[float,float,float] + Satellite velocity vector with respect to the primary body. + r_sec : tuple[float,float,float] Position vector of the secondary body with respect to the primary body. R_sec : float Equatorial radius of the secondary body. @@ -43,28 +45,27 @@ def eclipse_function(k, u_, r_sec, R_sec, R_primary, umbra=True): The current implementation assumes circular bodies and doesn't account for flattening. """ + # Plus or minus condition pm = 1 if umbra else -1 - p, ecc, inc, raan, argp, nu = rv2coe_hf( - k, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), RV2COE_TOL - ) + p, ecc, inc, raan, argp, nu = rv2coe_hf(k, rr, vv, RV2COE_TOL) - PQW = np.array(coe_rotation_matrix_hf(inc, raan, argp)) - # Make arrays contiguous for faster dot product with numba. - P_, Q_ = np.ascontiguousarray(PQW[:, 0]), np.ascontiguousarray(PQW[:, 1]) + PQW = coe_rotation_matrix_hf(inc, raan, argp) + P_ = PQW[0][0], PQW[1][0], PQW[2][0] + Q_ = PQW[0][1], PQW[1][1], PQW[2][1] - r_sec_norm = norm_V_hf(array_to_V_hf(r_sec)) - beta = (P_ @ r_sec) / r_sec_norm - zeta = (Q_ @ r_sec) / r_sec_norm + r_sec_norm = norm_V_hf(r_sec) + beta = matmul_VV_hf(P_, r_sec) / r_sec_norm + zeta = matmul_VV_hf(Q_, r_sec) / r_sec_norm - sin_delta_shadow = np.sin((R_sec - pm * R_primary) / r_sec_norm) + sin_delta_shadow = sin((R_sec - pm * R_primary) / r_sec_norm) - cos_psi = beta * np.cos(nu) + zeta * np.sin(nu) + cos_psi = beta * cos(nu) + zeta * sin(nu) shadow_function = ( - ((R_primary**2) * (1 + ecc * np.cos(nu)) ** 2) + ((R_primary**2) * (1 + ecc * cos(nu)) ** 2) + (p**2) * (cos_psi**2) - p**2 - + pm * (2 * p * R_primary * cos_psi) * (1 + ecc * np.cos(nu)) * sin_delta_shadow + + pm * (2 * p * R_primary * cos_psi) * (1 + ecc * cos(nu)) * sin_delta_shadow ) return shadow_function diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index f4bfd7a57..001d08d52 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -5,9 +5,10 @@ from astropy.coordinates import get_body_barycentric_posvel import numpy as np +from hapsira.core.jit import array_to_V_hf from hapsira.core.math.linalg import norm_V_vf from hapsira.core.events import ( - eclipse_function as eclipse_function_fast, + eclipse_function_hf, line_of_sight_gf, ) from hapsira.core.spheroid_location import ( @@ -199,14 +200,15 @@ def __call__(self, t, u_, k): self._last_t = t r_sec = super().__call__(t, u_, k) - shadow_function = eclipse_function_fast( + shadow_function = eclipse_function_hf( self.k, - u_, - r_sec, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), + array_to_V_hf(r_sec), self.R_sec, self.R_primary, - umbra=False, - ) + False, + ) # TODO call into hf return shadow_function @@ -233,8 +235,14 @@ def __call__(self, t, u_, k): self._last_t = t r_sec = super().__call__(t, u_, k) - shadow_function = eclipse_function_fast( - self.k, u_, r_sec, self.R_sec, self.R_primary + shadow_function = eclipse_function_hf( + self.k, + array_to_V_hf(u_[:3]), + array_to_V_hf(u_[3:]), + array_to_V_hf(r_sec), + self.R_sec, + self.R_primary, + True, ) return shadow_function From f7745ee8589d66aeb0fa030bdd9110380eea6bf7 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 20:49:16 +0100 Subject: [PATCH 206/346] jit cartesian_to_ellipsoidal --- src/hapsira/core/spheroid_location.py | 44 +++++++++++++++++++++------ src/hapsira/spheroid_location.py | 10 ++++-- src/hapsira/twobody/events.py | 8 ++--- 3 files changed, 46 insertions(+), 16 deletions(-) diff --git a/src/hapsira/core/spheroid_location.py b/src/hapsira/core/spheroid_location.py index c92ff1095..ef21f347e 100644 --- a/src/hapsira/core/spheroid_location.py +++ b/src/hapsira/core/spheroid_location.py @@ -1,12 +1,27 @@ """Low level calculations for oblate spheroid locations.""" +from math import atan, atan2, cos, sin, sqrt + from numba import njit as jit import numpy as np -from .jit import array_to_V_hf +from .jit import array_to_V_hf, hjit, gjit from .math.linalg import norm_V_hf +__all__ = [ + "cartesian_cords", + "f", + "N", + "tangential_vecs", + "radius_of_curvature", + "distance", + "is_visible", + "cartesian_to_ellipsoidal_hf", + "cartesian_to_ellipsoidal_gf", +] + + @jit def cartesian_cords(a, c, lon, lat, h): """Calculates cartesian coordinates. @@ -156,8 +171,8 @@ def is_visible(cartesian_cords, px, py, pz, N): return p >= 0 -@jit -def cartesian_to_ellipsoidal(a, c, x, y, z): +@hjit("Tuple([f,f,f])(f,f,f,f,f)") +def cartesian_to_ellipsoidal_hf(a, c, x, y, z): """Converts cartesian coordinates to ellipsoidal coordinates for the given ellipsoid. Instead of the iterative formula, the function uses the approximation introduced in Bowring, B. R. (1976). TRANSFORMATION FROM SPATIAL TO GEOGRAPHICAL COORDINATES. @@ -178,16 +193,25 @@ def cartesian_to_ellipsoidal(a, c, x, y, z): """ e2 = 1 - (c / a) ** 2 e2_ = e2 / (1 - e2) - p = np.sqrt(x**2 + y**2) - th = np.arctan(z * a / (p * c)) - lon = np.arctan2(y, x) # Use `arctan2` so that lon lies in the range: [-pi, +pi] - lat = np.arctan((z + e2_ * c * np.sin(th) ** 3) / (p - e2 * a * np.cos(th) ** 3)) + p = sqrt(x**2 + y**2) + th = atan(z * a / (p * c)) + lon = atan2(y, x) # Use `atan2` so that lon lies in the range: [-pi, +pi] + lat = atan((z + e2_ * c * sin(th) ** 3) / (p - e2 * a * cos(th) ** 3)) - v = a / np.sqrt(1 - e2 * np.sin(lat) ** 2) + v = a / sqrt(1 - e2 * sin(lat) ** 2) h = ( - np.sqrt(x**2 + y**2) / np.cos(lat) - v + sqrt(x**2 + y**2) / cos(lat) - v if lat < abs(1e-18) # to avoid errors very close and at zero - else z / np.sin(lat) - (1 - e2) * v + else z / sin(lat) - (1 - e2) * v ) return lon, lat, h + + +@gjit("void(f,f,f,f,f,f[:],f[:],f[:])", "(),(),(),(),()->(),(),()") +def cartesian_to_ellipsoidal_gf(a, c, x, y, z, lon, lat, h): + """ + Vectorized cartesian_to_ellipsoidal + """ + + lon[0], lat[0], h[0] = cartesian_to_ellipsoidal_hf(a, c, x, y, z) diff --git a/src/hapsira/spheroid_location.py b/src/hapsira/spheroid_location.py index c1a9f5073..e466f6258 100644 --- a/src/hapsira/spheroid_location.py +++ b/src/hapsira/spheroid_location.py @@ -4,7 +4,7 @@ from hapsira.core.spheroid_location import ( N as N_fast, cartesian_cords as cartesian_cords_fast, - cartesian_to_ellipsoidal as cartesian_to_ellipsoidal_fast, + cartesian_to_ellipsoidal_gf, distance as distance_fast, f as f_fast, is_visible as is_visible_fast, @@ -145,5 +145,11 @@ def cartesian_to_ellipsoidal(self, x, y, z): """ _a, _c = self._a.to_value(u.m), self._c.to_value(u.m) x, y, z = x.to_value(u.m), y.to_value(u.m), z.to_value(u.m) - lon, lat, h = cartesian_to_ellipsoidal_fast(_a, _c, x, y, z) + lon, lat, h = cartesian_to_ellipsoidal_gf( # pylint: disable=E1120,E0633 + _a, + _c, + x, + y, + z, + ) return lon * u.rad, lat * u.rad, h * u.m diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index 001d08d52..0b708e868 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -11,9 +11,7 @@ eclipse_function_hf, line_of_sight_gf, ) -from hapsira.core.spheroid_location import ( - cartesian_to_ellipsoidal as cartesian_to_ellipsoidal_fast, -) +from hapsira.core.spheroid_location import cartesian_to_ellipsoidal_hf __all__ = [ @@ -138,7 +136,9 @@ def __init__(self, orbit, lat, terminal=False, direction=0): def __call__(self, t, u_, k): self._last_t = t pos_on_body = (u_[:3] / norm_V_vf(*u_[:3])) * self._R - _, lat_, _ = cartesian_to_ellipsoidal_fast(self._R, self._R_polar, *pos_on_body) + _, lat_, _ = cartesian_to_ellipsoidal_hf( + self._R, self._R_polar, *pos_on_body + ) # TODO call into hf return np.rad2deg(lat_) - self._lat From 96affb014f30cf2268cd3f14a9e34275e0b0b1b1 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 21:42:31 +0100 Subject: [PATCH 207/346] cleanup --- src/hapsira/twobody/events.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index 0b708e868..edefd1d01 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -19,7 +19,7 @@ "AltitudeCrossEvent", "LithobrakeEvent", "LatitudeCrossEvent", - "EclipseEvent", + "BaseEclipseEvent", "PenumbraEvent", "UmbraEvent", "NodeCrossEvent", @@ -143,7 +143,7 @@ def __call__(self, t, u_, k): return np.rad2deg(lat_) - self._lat -class EclipseEvent(BaseEvent): +class BaseEclipseEvent(BaseEvent): """Base class for the eclipse event. Parameters @@ -166,6 +166,7 @@ def __init__(self, orbit, terminal=False, direction=0): self.R_sec = self._secondary_body.R.to_value(u.km) self.R_primary = self._primary_body.R.to_value(u.km) + @abstractmethod def __call__(self, t, u_, k): # Solve for primary and secondary bodies position w.r.t. solar system # barycenter at a particular epoch. @@ -178,7 +179,7 @@ def __call__(self, t, u_, k): return r_sec -class PenumbraEvent(EclipseEvent): +class PenumbraEvent(BaseEclipseEvent): """Detect whether a satellite is in penumbra or not. Parameters @@ -193,9 +194,6 @@ class PenumbraEvent(EclipseEvent): """ - def __init__(self, orbit, terminal=False, direction=0): - super().__init__(orbit, terminal, direction) - def __call__(self, t, u_, k): self._last_t = t @@ -213,7 +211,7 @@ def __call__(self, t, u_, k): return shadow_function -class UmbraEvent(EclipseEvent): +class UmbraEvent(BaseEclipseEvent): """Detect whether a satellite is in umbra or not. Parameters @@ -228,9 +226,6 @@ class UmbraEvent(EclipseEvent): """ - def __init__(self, orbit, terminal=False, direction=0): - super().__init__(orbit, terminal, direction) - def __call__(self, t, u_, k): self._last_t = t From 08c51e9cdc6d43f1b774ee0abc3e4e3cbfc27cc2 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 22:45:40 +0100 Subject: [PATCH 208/346] evaluating umbra events is based on interpolation vs calling get_body_barycentric_posvel --- src/hapsira/twobody/events.py | 32 ++++++++++++++++++++---------- tests/tests_twobody/test_events.py | 14 ++++++++----- 2 files changed, 31 insertions(+), 15 deletions(-) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index edefd1d01..f5dc93792 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -11,7 +11,9 @@ eclipse_function_hf, line_of_sight_gf, ) +from hapsira.core.math.interpolate import interp_hb from hapsira.core.spheroid_location import cartesian_to_ellipsoidal_hf +from hapsira.util import time_range __all__ = [ @@ -150,6 +152,10 @@ class BaseEclipseEvent(BaseEvent): ---------- orbit: hapsira.twobody.orbit.Orbit Orbit of the satellite. + tof: ~astropy.units.Quantity + Maximum time of flight for interpolator + steps: int + Steps for interpolator terminal: bool, optional Whether to terminate integration when the event occurs, defaults to False. direction: float, optional @@ -157,7 +163,7 @@ class BaseEclipseEvent(BaseEvent): """ - def __init__(self, orbit, terminal=False, direction=0): + def __init__(self, orbit, tof, steps=50, terminal=False, direction=0): super().__init__(terminal, direction) self._primary_body = orbit.attractor self._secondary_body = orbit.attractor.parent @@ -166,17 +172,23 @@ def __init__(self, orbit, terminal=False, direction=0): self.R_sec = self._secondary_body.R.to_value(u.km) self.R_primary = self._primary_body.R.to_value(u.km) + epochs = time_range(start=self._epoch, end=self._epoch + tof, num_values=steps) + r_primary_wrt_ssb, _ = get_body_barycentric_posvel( + self._primary_body.name, epochs + ) + r_secondary_wrt_ssb, _ = get_body_barycentric_posvel( + self._secondary_body.name, epochs + ) + self._r_sec = interp_hb( + (epochs - self._epoch).to_value(u.s), + (r_secondary_wrt_ssb - r_primary_wrt_ssb).xyz.to_value(u.km), + ) + @abstractmethod def __call__(self, t, u_, k): # Solve for primary and secondary bodies position w.r.t. solar system # barycenter at a particular epoch. - (r_primary_wrt_ssb, _), (r_secondary_wrt_ssb, _) = ( - get_body_barycentric_posvel(body.name, self._epoch + t * u.s) - for body in (self._primary_body, self._secondary_body) - ) - r_sec = ((r_secondary_wrt_ssb - r_primary_wrt_ssb).xyz << u.km).value - - return r_sec + return self._r_sec(t) class PenumbraEvent(BaseEclipseEvent): @@ -202,7 +214,7 @@ def __call__(self, t, u_, k): self.k, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), - array_to_V_hf(r_sec), + r_sec, self.R_sec, self.R_primary, False, @@ -234,7 +246,7 @@ def __call__(self, t, u_, k): self.k, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), - array_to_V_hf(r_sec), + r_sec, self.R_sec, self.R_primary, True, diff --git a/tests/tests_twobody/test_events.py b/tests/tests_twobody/test_events.py index 1b0a37dc9..09db9d8aa 100644 --- a/tests/tests_twobody/test_events.py +++ b/tests/tests_twobody/test_events.py @@ -122,7 +122,7 @@ def test_penumbra_event_not_triggering_is_ok(): v0 = np.array([7.36138, 2.98997, 1.64354]) orbit = Orbit.from_vectors(attractor, r0 * u.km, v0 * u.km / u.s) - penumbra_event = PenumbraEvent(orbit) + penumbra_event = PenumbraEvent(orbit, tof=tof) method = CowellPropagator(events=[penumbra_event]) rr, _ = method.propagate_many( orbit._state, @@ -139,7 +139,7 @@ def test_umbra_event_not_triggering_is_ok(): v0 = np.array([7.36138, 2.98997, 1.64354]) orbit = Orbit.from_vectors(attractor, r0 * u.km, v0 * u.km / u.s) - umbra_event = UmbraEvent(orbit) + umbra_event = UmbraEvent(orbit, tof=tof) method = CowellPropagator(events=[umbra_event]) rr, _ = method.propagate_many( @@ -166,7 +166,7 @@ def test_umbra_event_crossing(): epoch=epoch, ) - umbra_event = UmbraEvent(orbit, terminal=True) + umbra_event = UmbraEvent(orbit, tof=tof, terminal=True) method = CowellPropagator(events=[umbra_event]) rr, _ = method.propagate_many( @@ -193,7 +193,7 @@ def test_penumbra_event_crossing(): epoch=epoch, ) - penumbra_event = PenumbraEvent(orbit, terminal=True) + penumbra_event = PenumbraEvent(orbit, tof=tof, terminal=True) method = CowellPropagator(events=[penumbra_event]) rr, _ = method.propagate_many( orbit._state, @@ -310,7 +310,11 @@ def test_propagation_stops_if_atleast_one_event_has_terminal_set_to_True( epoch=epoch, ) - penumbra_event = PenumbraEvent(orbit, terminal=penumbra_terminal) + penumbra_event = PenumbraEvent( + orbit, + tof=600 * u.s, + terminal=penumbra_terminal, + ) thresh_lat = 30 * u.deg latitude_cross_event = LatitudeCrossEvent( From 47e43f5972fefa83dc3cd6a5c70c85053310aa66 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Jan 2024 22:47:11 +0100 Subject: [PATCH 209/346] fix name --- src/hapsira/twobody/events.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index f5dc93792..8328a8daa 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -179,7 +179,7 @@ def __init__(self, orbit, tof, steps=50, terminal=False, direction=0): r_secondary_wrt_ssb, _ = get_body_barycentric_posvel( self._secondary_body.name, epochs ) - self._r_sec = interp_hb( + self._r_sec_hf = interp_hb( (epochs - self._epoch).to_value(u.s), (r_secondary_wrt_ssb - r_primary_wrt_ssb).xyz.to_value(u.km), ) @@ -188,7 +188,7 @@ def __init__(self, orbit, tof, steps=50, terminal=False, direction=0): def __call__(self, t, u_, k): # Solve for primary and secondary bodies position w.r.t. solar system # barycenter at a particular epoch. - return self._r_sec(t) + return self._r_sec_hf(t) class PenumbraEvent(BaseEclipseEvent): From 988a91cd22ff1de285c33af35b13d288963fcaf7 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 28 Jan 2024 12:42:06 +0100 Subject: [PATCH 210/346] los event runs on 1D linear interpolation function for secondary body positions --- src/hapsira/twobody/events.py | 12 +++--------- tests/tests_twobody/test_events.py | 29 +++++++++++++---------------- 2 files changed, 16 insertions(+), 25 deletions(-) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index 8328a8daa..8e87de4c5 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -291,13 +291,10 @@ class LosEvent(BaseEvent): """ - def __init__(self, attractor, pos_coords, terminal=False, direction=0): + def __init__(self, attractor, tofs, secondary_rr, terminal=False, direction=0): super().__init__(terminal, direction) self._attractor = attractor - self._pos_coords = (pos_coords << u.km).value.tolist() - self._last_coord = ( - self._pos_coords[-1] << u.km - ).value # Used to prevent any errors if `self._pos_coords` gets exhausted early. + self._secondary_hf = interp_hb(tofs.to_value(u.s), secondary_rr.to_value(u.km)) self._R = self._attractor.R.to_value(u.km) def __call__(self, t, u_, k): @@ -308,12 +305,9 @@ def __call__(self, t, u_, k): "The norm of the position vector of the primary body is less than the radius of the attractor." ) - pos_coord = self._pos_coords.pop(0) if self._pos_coords else self._last_coord - - # Need to cast `pos_coord` to array since `norm` inside numba only works for arrays, not lists. delta_angle = line_of_sight_gf( # pylint: disable=E1120 u_[:3], - np.array(pos_coord), + np.array(self._secondary_hf(t)), self._R, ) return delta_angle diff --git a/tests/tests_twobody/test_events.py b/tests/tests_twobody/test_events.py index 09db9d8aa..c3f420de6 100644 --- a/tests/tests_twobody/test_events.py +++ b/tests/tests_twobody/test_events.py @@ -356,14 +356,13 @@ def test_LOS_event_raises_warning_if_norm_of_r1_less_than_attractor_radius_durin v2 = np.array([5021.38, -2900.7, 1000.354]) << u.km / u.s orbit = Orbit.from_vectors(Earth, r2, v2) - tofs = [100, 500, 1000, 2000] << u.s + tofs = tofs = np.arange(0, 2000, 10) << u.s # Propagate the secondary body to generate its position coordinates method = CowellPropagator() - rr, vv = method.propagate_many( + secondary_rr, _ = method.propagate_many( orbit._state, tofs, ) - pos_coords = rr # Trajectory of the secondary body. r1 = ( np.array([0, -5010.696, -5102.509]) << u.km @@ -371,7 +370,7 @@ def test_LOS_event_raises_warning_if_norm_of_r1_less_than_attractor_radius_durin v1 = np.array([736.138, 29899.7, 164.354]) << u.km / u.s orb = Orbit.from_vectors(Earth, r1, v1) - los_event = LosEvent(Earth, pos_coords, terminal=True) + los_event = LosEvent(Earth, tofs, secondary_rr.T, terminal=True) events = [los_event] tofs = [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.5] << u.s @@ -389,20 +388,19 @@ def test_LOS_event_with_lithobrake_event_raises_warning_when_satellite_cuts_attr v2 = np.array([5021.38, -2900.7, 1000.354]) << u.km / u.s orbit = Orbit.from_vectors(Earth, r2, v2) - tofs = [100, 500, 1000, 2000] << u.s + tofs = tofs = np.arange(0, 2000, 10) << u.s # Propagate the secondary body to generate its position coordinates method = CowellPropagator() - rr, vv = method.propagate_many( + secondary_rr, _ = method.propagate_many( orbit._state, tofs, ) - pos_coords = rr # Trajectory of the secondary body. r1 = np.array([0, -5010.696, -5102.509]) << u.km v1 = np.array([736.138, 2989.7, 164.354]) << u.km / u.s orb = Orbit.from_vectors(Earth, r1, v1) - los_event = LosEvent(Earth, pos_coords, terminal=True) + los_event = LosEvent(Earth, tofs, secondary_rr.T, terminal=True) tofs = [ 0.003, 0.004, @@ -421,7 +419,7 @@ def test_LOS_event_with_lithobrake_event_raises_warning_when_satellite_cuts_attr lithobrake_event = LithobrakeEvent(Earth.R.to_value(u.km)) method = CowellPropagator(events=[lithobrake_event, los_event]) - r, v = method.propagate_many( + _, _ = method.propagate_many( orb._state, tofs, ) @@ -430,19 +428,19 @@ def test_LOS_event_with_lithobrake_event_raises_warning_when_satellite_cuts_attr def test_LOS_event(): - t_los = 2327.165 * u.s + t_los = 2327.381434 * u.s r2 = np.array([-500, 1500, 4012.09]) << u.km v2 = np.array([5021.38, -2900.7, 1000.354]) << u.km / u.s orbit = Orbit.from_vectors(Earth, r2, v2) - tofs = [100, 500, 1000, 2000] << u.s + tofs = np.arange(0, 5000, 10) << u.s + # Propagate the secondary body to generate its position coordinates method = CowellPropagator() - rr, vv = method.propagate_many( + secondary_rr, _ = method.propagate_many( orbit._state, tofs, ) - pos_coords = rr # Trajectory of the secondary body. orb = Orbit.from_classical( attractor=Earth, @@ -454,12 +452,11 @@ def test_LOS_event(): nu=30 * u.deg, ) - los_event = LosEvent(Earth, pos_coords, terminal=True) + los_event = LosEvent(Earth, tofs, secondary_rr.T, terminal=True) events = [los_event] - tofs = [1, 5, 10, 100, 1000, 2000, 3000, 5000] << u.s method = CowellPropagator(events=events) - r, v = method.propagate_many( + _, _ = method.propagate_many( orb._state, tofs, ) From b2ba64ae5eadc23fd57d200e45727d5df3dd2b6d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 28 Jan 2024 14:32:55 +0100 Subject: [PATCH 211/346] events get compiled --- src/hapsira/twobody/events.py | 183 +++++++++++++++++++++------------- 1 file changed, 113 insertions(+), 70 deletions(-) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index 8e87de4c5..ebb6e5bc6 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -1,15 +1,15 @@ from abc import ABC, abstractmethod -from warnings import warn +from math import degrees as rad2deg +from typing import Callable from astropy import units as u from astropy.coordinates import get_body_barycentric_posvel -import numpy as np -from hapsira.core.jit import array_to_V_hf -from hapsira.core.math.linalg import norm_V_vf +from hapsira.core.jit import hjit +from hapsira.core.math.linalg import mul_Vs_hf, norm_V_hf from hapsira.core.events import ( eclipse_function_hf, - line_of_sight_gf, + line_of_sight_hf, ) from hapsira.core.math.interpolate import interp_hb from hapsira.core.spheroid_location import cartesian_to_ellipsoidal_hf @@ -41,9 +41,11 @@ class BaseEvent(ABC): """ + @abstractmethod def __init__(self, terminal, direction): self._terminal, self._direction = terminal, direction self._last_t = None + self._impl_hf = None @property def terminal(self): @@ -57,9 +59,21 @@ def direction(self): def last_t(self): return self._last_t << u.s - @abstractmethod - def __call__(self, t, u, k): - raise NotImplementedError + @property + def last_t_raw(self) -> float: + return self._last_t + + @last_t_raw.setter + def last_t_raw(self, value: float): + self._last_t = value + + @property + def impl_hf(self) -> Callable: + return self._impl_hf + + def __call__(self, t, rr, vv, k): + self._last_t = t + return self._impl_hf(t, rr, vv, k) class AltitudeCrossEvent(BaseEvent): @@ -85,13 +99,14 @@ def __init__(self, alt, R, terminal=True, direction=-1): self._R = R self._alt = alt # Threshold altitude from the ground. - def __call__(self, t, u, k): - self._last_t = t - r_norm = norm_V_vf(*u[:3]) + @hjit("f(f,V,V,f)", cache=False) + def impl_hf(t, rr, vv, k): + r_norm = norm_V_hf(rr) + return ( + r_norm - R - alt + ) # If this goes from +ve to -ve, altitude is decreasing. - return ( - r_norm - self._R - self._alt - ) # If this goes from +ve to -ve, altitude is decreasing. + self._impl_hf = impl_hf class LithobrakeEvent(AltitudeCrossEvent): @@ -135,14 +150,17 @@ def __init__(self, orbit, lat, terminal=False, direction=0): self._epoch = orbit.epoch self._lat = lat.to_value(u.deg) # Threshold latitude (in degrees). - def __call__(self, t, u_, k): - self._last_t = t - pos_on_body = (u_[:3] / norm_V_vf(*u_[:3])) * self._R - _, lat_, _ = cartesian_to_ellipsoidal_hf( - self._R, self._R_polar, *pos_on_body - ) # TODO call into hf + _R = self._R + _R_polar = self._R_polar + _lat = self._lat + + @hjit("f(f,V,V,f)", cache=False) + def impl_hf(t, rr, vv, k): + pos_on_body = mul_Vs_hf(rr, _R / norm_V_hf(rr)) + _, lat_, _ = cartesian_to_ellipsoidal_hf(_R, _R_polar, *pos_on_body) + return rad2deg(lat_) - _lat - return np.rad2deg(lat_) - self._lat + self._impl_hf = impl_hf class BaseEclipseEvent(BaseEvent): @@ -184,11 +202,15 @@ def __init__(self, orbit, tof, steps=50, terminal=False, direction=0): (r_secondary_wrt_ssb - r_primary_wrt_ssb).xyz.to_value(u.km), ) - @abstractmethod - def __call__(self, t, u_, k): - # Solve for primary and secondary bodies position w.r.t. solar system - # barycenter at a particular epoch. - return self._r_sec_hf(t) + _r_sec_hf = self._r_sec_hf + + @hjit("V(f,V,V,f)", cache=False) + def impl_hf(t, rr, vv, k): + # Solve for primary and secondary bodies position w.r.t. + # solar system barycenter at a particular epoch. + return _r_sec_hf(t) + + self._impl_hf = impl_hf class PenumbraEvent(BaseEclipseEvent): @@ -206,21 +228,28 @@ class PenumbraEvent(BaseEclipseEvent): """ - def __call__(self, t, u_, k): - self._last_t = t - - r_sec = super().__call__(t, u_, k) - shadow_function = eclipse_function_hf( - self.k, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), - r_sec, - self.R_sec, - self.R_primary, - False, - ) # TODO call into hf + def __init__(self, orbit, tof, steps=50, terminal=False, direction=0): + super().__init__(orbit, tof, steps, terminal, direction) + + R_sec = self.R_sec + R_primary = self.R_primary + _impl_hf = self._impl_hf + + @hjit("f(f,V,V,f)", cache=False) + def impl_hf(t, rr, vv, k): + r_sec = _impl_hf(t, rr, vv, k) + shadow_function = eclipse_function_hf( + k, + rr, + vv, + r_sec, + R_sec, + R_primary, + False, + ) + return shadow_function - return shadow_function + self._impl_hf = impl_hf class UmbraEvent(BaseEclipseEvent): @@ -238,21 +267,28 @@ class UmbraEvent(BaseEclipseEvent): """ - def __call__(self, t, u_, k): - self._last_t = t - - r_sec = super().__call__(t, u_, k) - shadow_function = eclipse_function_hf( - self.k, - array_to_V_hf(u_[:3]), - array_to_V_hf(u_[3:]), - r_sec, - self.R_sec, - self.R_primary, - True, - ) + def __init__(self, orbit, tof, steps=50, terminal=False, direction=0): + super().__init__(orbit, tof, steps, terminal, direction) + + R_sec = self.R_sec + R_primary = self.R_primary + _impl_hf = self._impl_hf + + @hjit("f(f,V,V,f)", cache=False) + def impl_hf(t, rr, vv, k): + r_sec = _impl_hf(t, rr, vv, k) + shadow_function = eclipse_function_hf( + k, + rr, + vv, + r_sec, + R_sec, + R_primary, + True, + ) + return shadow_function - return shadow_function + self._impl_hf = impl_hf class NodeCrossEvent(BaseEvent): @@ -272,10 +308,12 @@ class NodeCrossEvent(BaseEvent): def __init__(self, terminal=False, direction=0): super().__init__(terminal, direction) - def __call__(self, t, u_, k): - self._last_t = t - # Check if the z coordinate of the satellite is zero. - return u_[2] + @hjit("f(f,V,V,f)", cache=False) + def impl_hf(t, rr, vv, k): + # Check if the z coordinate of the satellite is zero. + return rr[2] + + self._impl_hf = impl_hf class LosEvent(BaseEvent): @@ -297,17 +335,22 @@ def __init__(self, attractor, tofs, secondary_rr, terminal=False, direction=0): self._secondary_hf = interp_hb(tofs.to_value(u.s), secondary_rr.to_value(u.km)) self._R = self._attractor.R.to_value(u.km) - def __call__(self, t, u_, k): - self._last_t = t - - if norm_V_vf(*u_[:3]) < self._R: - warn( - "The norm of the position vector of the primary body is less than the radius of the attractor." + _R = self._R + _secondary_hf = self._secondary_hf + + @hjit("f(f,V,V,f)", cache=False) + def impl_hf(t, rr, vv, k): + # Can currently not warn due to: https://github.com/numba/numba/issues/1243 + # TODO Matching test deactivated ... + # if norm_V_hf(rr) < _R: + # warn( + # "The norm of the position vector of the primary body is less than the radius of the attractor." + # ) + delta_angle = line_of_sight_hf( + rr, + _secondary_hf(t), + _R, ) + return delta_angle - delta_angle = line_of_sight_gf( # pylint: disable=E1120 - u_[:3], - np.array(self._secondary_hf(t)), - self._R, - ) - return delta_angle + self._impl_hf = impl_hf From 62f3ef58ac80378b0318fdac94f6ecc0ee003faf Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 28 Jan 2024 14:33:23 +0100 Subject: [PATCH 212/346] deactivate test for warning as compiled code currently can not raise them --- tests/tests_twobody/test_events.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/tests_twobody/test_events.py b/tests/tests_twobody/test_events.py index c3f420de6..af2b9c7ff 100644 --- a/tests/tests_twobody/test_events.py +++ b/tests/tests_twobody/test_events.py @@ -374,12 +374,14 @@ def test_LOS_event_raises_warning_if_norm_of_r1_less_than_attractor_radius_durin events = [los_event] tofs = [0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.5] << u.s - with pytest.warns(UserWarning, match="The norm of the position vector"): - method = CowellPropagator(events=events) - r, v = method.propagate_many( - orb._state, - tofs, - ) + # Can currently not warn due to: https://github.com/numba/numba/issues/1243 + # TODO Matching implementation in LosEvent deactivated + # with pytest.warns(UserWarning, match="The norm of the position vector"): + method = CowellPropagator(events=events) + _, _ = method.propagate_many( + orb._state, + tofs, + ) # should trigger waring @pytest.mark.filterwarnings("ignore::UserWarning") From 2978c398ab363d808c0cd527c96103b73064b4e4 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 28 Jan 2024 14:34:24 +0100 Subject: [PATCH 213/346] less arrays, more tuples --- src/hapsira/core/math/ivp/_rk.py | 3 ++- src/hapsira/core/math/ivp/_solve.py | 9 ++++----- src/hapsira/core/propagation/cowell.py | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 73d5f7ac4..413c27017 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -234,7 +234,8 @@ def __call__(self, t): y *= 1 - x y += self.y_old - return y.T + y = y.reshape(6) + return array_to_V_hf(y[:3]), array_to_V_hf(y[3:]) class DOP853: diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index b35e8aced..59aa6f76a 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -40,7 +40,8 @@ def _solve_event_equation( """ def wrapper(t): - return event(t, sol(t), argk) + rr, vv = sol(t) + return event(t, rr, vv, argk) value, status = brentq_hf( wrapper, @@ -203,7 +204,6 @@ def solve_ivp( solver = DOP853(fun, t0, rr, vv, tf, argk, rtol, atol) - y0 = np.array([*rr, *vv]) # TODO turn into tuples ts = [t0] interpolants = [] @@ -211,7 +211,7 @@ def solve_ivp( events, is_terminal, event_dir = _prepare_events(events) if events is not None: - g = [event(t0, y0, argk) for event in events] + g = [event(t0, rr, vv, argk) for event in events] status = None while status is None: @@ -225,13 +225,12 @@ def solve_ivp( t_old = solver.t_old t = solver.t - y = np.array([*solver.rr, *solver.vv]) # TODO turn into tuples sol = solver.dense_output() interpolants.append(sol) if events is not None: - g_new = [event(t, y, argk) for event in events] + g_new = [event(t, solver.rr, solver.vv, argk) for event in events] active_events = _find_active_events(g, g_new, event_dir) if active_events.size > 0: _, roots, terminate = _handle_events( diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 5766fb7a6..b0f2c96a3 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -55,8 +55,8 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=None, f=func_twobody_hf rrs = [] vvs = [] for t in tofs: - y = sol(t) - rrs.append(y[:3]) - vvs.append(y[3:]) + r_new, v_new = sol(t) + rrs.append(r_new) + vvs.append(v_new) return rrs, vvs From c5cd8e69e29cbf98af299c6bd1d62d54e327d883 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 28 Jan 2024 14:58:10 +0100 Subject: [PATCH 214/346] cleanup --- src/hapsira/twobody/events.py | 85 +++++++++++++---------------------- 1 file changed, 31 insertions(+), 54 deletions(-) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index ebb6e5bc6..d9bb60e09 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -82,7 +82,7 @@ class AltitudeCrossEvent(BaseEvent): Parameters ---------- alt: float - Threshold altitude (km). + Threshold altitude from the ground (km). R: float Radius of the attractor (km). terminal: bool @@ -96,8 +96,6 @@ class AltitudeCrossEvent(BaseEvent): def __init__(self, alt, R, terminal=True, direction=-1): super().__init__(terminal, direction) - self._R = R - self._alt = alt # Threshold altitude from the ground. @hjit("f(f,V,V,f)", cache=False) def impl_hf(t, rr, vv, k): @@ -145,20 +143,15 @@ class LatitudeCrossEvent(BaseEvent): def __init__(self, orbit, lat, terminal=False, direction=0): super().__init__(terminal, direction) - self._R = orbit.attractor.R.to_value(u.m) - self._R_polar = orbit.attractor.R_polar.to_value(u.m) - self._epoch = orbit.epoch - self._lat = lat.to_value(u.deg) # Threshold latitude (in degrees). - - _R = self._R - _R_polar = self._R_polar - _lat = self._lat + R = orbit.attractor.R.to_value(u.m) + R_polar = orbit.attractor.R_polar.to_value(u.m) + lat = lat.to_value(u.deg) # Threshold latitude (in degrees). @hjit("f(f,V,V,f)", cache=False) def impl_hf(t, rr, vv, k): - pos_on_body = mul_Vs_hf(rr, _R / norm_V_hf(rr)) - _, lat_, _ = cartesian_to_ellipsoidal_hf(_R, _R_polar, *pos_on_body) - return rad2deg(lat_) - _lat + pos_on_body = mul_Vs_hf(rr, R / norm_V_hf(rr)) + _, lat_, _ = cartesian_to_ellipsoidal_hf(R, R_polar, *pos_on_body) + return rad2deg(lat_) - lat self._impl_hf = impl_hf @@ -183,35 +176,23 @@ class BaseEclipseEvent(BaseEvent): def __init__(self, orbit, tof, steps=50, terminal=False, direction=0): super().__init__(terminal, direction) - self._primary_body = orbit.attractor - self._secondary_body = orbit.attractor.parent - self._epoch = orbit.epoch - self.k = self._primary_body.k.to_value(u.km**3 / u.s**2) - self.R_sec = self._secondary_body.R.to_value(u.km) - self.R_primary = self._primary_body.R.to_value(u.km) - - epochs = time_range(start=self._epoch, end=self._epoch + tof, num_values=steps) - r_primary_wrt_ssb, _ = get_body_barycentric_posvel( - self._primary_body.name, epochs - ) + primary_body = orbit.attractor + secondary_body = orbit.attractor.parent + epoch = orbit.epoch + + self._R_sec = secondary_body.R.to_value(u.km) + self._R_primary = primary_body.R.to_value(u.km) + + epochs = time_range(start=epoch, end=epoch + tof, num_values=steps) + r_primary_wrt_ssb, _ = get_body_barycentric_posvel(primary_body.name, epochs) r_secondary_wrt_ssb, _ = get_body_barycentric_posvel( - self._secondary_body.name, epochs + secondary_body.name, epochs ) self._r_sec_hf = interp_hb( - (epochs - self._epoch).to_value(u.s), + (epochs - epoch).to_value(u.s), (r_secondary_wrt_ssb - r_primary_wrt_ssb).xyz.to_value(u.km), ) - _r_sec_hf = self._r_sec_hf - - @hjit("V(f,V,V,f)", cache=False) - def impl_hf(t, rr, vv, k): - # Solve for primary and secondary bodies position w.r.t. - # solar system barycenter at a particular epoch. - return _r_sec_hf(t) - - self._impl_hf = impl_hf - class PenumbraEvent(BaseEclipseEvent): """Detect whether a satellite is in penumbra or not. @@ -231,13 +212,13 @@ class PenumbraEvent(BaseEclipseEvent): def __init__(self, orbit, tof, steps=50, terminal=False, direction=0): super().__init__(orbit, tof, steps, terminal, direction) - R_sec = self.R_sec - R_primary = self.R_primary - _impl_hf = self._impl_hf + R_sec = self._R_sec + R_primary = self._R_primary + r_sec_hf = self._r_sec_hf @hjit("f(f,V,V,f)", cache=False) def impl_hf(t, rr, vv, k): - r_sec = _impl_hf(t, rr, vv, k) + r_sec = r_sec_hf(t) shadow_function = eclipse_function_hf( k, rr, @@ -270,13 +251,13 @@ class UmbraEvent(BaseEclipseEvent): def __init__(self, orbit, tof, steps=50, terminal=False, direction=0): super().__init__(orbit, tof, steps, terminal, direction) - R_sec = self.R_sec - R_primary = self.R_primary - _impl_hf = self._impl_hf + R_sec = self._R_sec + R_primary = self._R_primary + r_sec_hf = self._r_sec_hf @hjit("f(f,V,V,f)", cache=False) def impl_hf(t, rr, vv, k): - r_sec = _impl_hf(t, rr, vv, k) + r_sec = r_sec_hf(t) shadow_function = eclipse_function_hf( k, rr, @@ -331,25 +312,21 @@ class LosEvent(BaseEvent): def __init__(self, attractor, tofs, secondary_rr, terminal=False, direction=0): super().__init__(terminal, direction) - self._attractor = attractor - self._secondary_hf = interp_hb(tofs.to_value(u.s), secondary_rr.to_value(u.km)) - self._R = self._attractor.R.to_value(u.km) - - _R = self._R - _secondary_hf = self._secondary_hf + secondary_hf = interp_hb(tofs.to_value(u.s), secondary_rr.to_value(u.km)) + R = attractor.R.to_value(u.km) @hjit("f(f,V,V,f)", cache=False) def impl_hf(t, rr, vv, k): # Can currently not warn due to: https://github.com/numba/numba/issues/1243 # TODO Matching test deactivated ... - # if norm_V_hf(rr) < _R: + # if norm_V_hf(rr) < R: # warn( # "The norm of the position vector of the primary body is less than the radius of the attractor." # ) delta_angle = line_of_sight_hf( rr, - _secondary_hf(t), - _R, + secondary_hf(t), + R, ) return delta_angle From 73cc39c01d8c7949243e94be8db59b9e91204a4c Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 28 Jan 2024 19:04:30 +0100 Subject: [PATCH 215/346] cleanup --- src/hapsira/core/math/ivp/_rk.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 413c27017..6d02ff8b2 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -194,9 +194,6 @@ class Dop853DenseOutput: def __init__(self, t_old, t, y_old, F): self.t_old = t_old - self.t = t - self.t_min = min(t, t_old) - self.t_max = max(t, t_old) self.h = t - t_old self.F = F self.y_old = y_old From b7cd3a1cc40d92f91d4f66e6608e2f6024d33de1 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 28 Jan 2024 20:29:37 +0100 Subject: [PATCH 216/346] prep for cleanup --- src/hapsira/core/math/ivp/_rk.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 6d02ff8b2..b10bf7b0f 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -198,7 +198,7 @@ def __init__(self, t_old, t, y_old, F): self.F = F self.y_old = y_old - def __call__(self, t): + def __call__(self, t: float): """Evaluate the interpolant. Parameters @@ -212,16 +212,12 @@ def __call__(self, t): Computed values. Shape depends on whether `t` was a scalar or a 1-D array. """ - t = np.asarray(t) - assert not t.ndim > 1 + # t = np.asarray(t) + # assert not t.ndim > 1 + assert np.asarray(t).shape == tuple() x = (t - self.t_old) / self.h - - if t.ndim == 0: - y = np.zeros_like(self.y_old) - else: - x = x[:, None] - y = np.zeros((len(x), len(self.y_old)), dtype=self.y_old.dtype) + y = np.zeros_like(self.y_old) for i, f in enumerate(reversed(self.F)): y += f @@ -231,7 +227,7 @@ def __call__(self, t): y *= 1 - x y += self.y_old - y = y.reshape(6) + assert y.shape == (6,) return array_to_V_hf(y[:3]), array_to_V_hf(y[3:]) From b0ecb22e986ac7c22adfe2a4d609dee8773449b0 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 28 Jan 2024 20:39:14 +0100 Subject: [PATCH 217/346] cleanup dense output so it relies on 4 unchanged states --- src/hapsira/core/math/ivp/_rk.py | 40 +++++++++++++++++++------------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index b10bf7b0f..25ba51356 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -192,11 +192,12 @@ class Dop853DenseOutput: Time range of the interpolation. """ - def __init__(self, t_old, t, y_old, F): + def __init__(self, t_old, h, rr_old, vv_old, F): self.t_old = t_old - self.h = t - t_old + self.h = h self.F = F - self.y_old = y_old + self.rr_old = rr_old + self.vv_old = vv_old def __call__(self, t: float): """Evaluate the interpolant. @@ -212,23 +213,26 @@ def __call__(self, t: float): Computed values. Shape depends on whether `t` was a scalar or a 1-D array. """ - # t = np.asarray(t) - # assert not t.ndim > 1 - assert np.asarray(t).shape == tuple() x = (t - self.t_old) / self.h - y = np.zeros_like(self.y_old) + rr_new = (0.0, 0.0, 0.0) + vv_new = (0.0, 0.0, 0.0) - for i, f in enumerate(reversed(self.F)): - y += f - if i % 2 == 0: - y *= x + for idx, f in enumerate(reversed(self.F)): + rr_new = add_VV_hf(rr_new, array_to_V_hf(f[:3])) + vv_new = add_VV_hf(vv_new, array_to_V_hf(f[3:])) + + if idx % 2 == 0: + rr_new = mul_Vs_hf(rr_new, x) + vv_new = mul_Vs_hf(vv_new, x) else: - y *= 1 - x - y += self.y_old + rr_new = mul_Vs_hf(rr_new, 1 - x) + vv_new = mul_Vs_hf(vv_new, 1 - x) + + rr_new = add_VV_hf(rr_new, self.rr_old) + vv_new = add_VV_hf(vv_new, self.vv_old) - assert y.shape == (6,) - return array_to_V_hf(y[:3]), array_to_V_hf(y[3:]) + return rr_new, vv_new class DOP853: @@ -408,5 +412,9 @@ def dense_output(self): F[3:, :] = h * np.dot(self.D, K) # TODO return Dop853DenseOutput( - self.t_old, self.t, np.array([*self.rr_old, *self.vv_old]), F + self.t_old, + self.t - self.t_old, # h + self.rr_old, + self.vv_old, + F, ) From 96f39089decd32362c854b7c7e96dbacc40e5459 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 28 Jan 2024 21:03:06 +0100 Subject: [PATCH 218/346] dense output entirely on tuples --- src/hapsira/core/math/ivp/_rk.py | 36 ++++++++++++++------------------ 1 file changed, 16 insertions(+), 20 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 25ba51356..a8e3430b1 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -10,7 +10,6 @@ SAFETY, MIN_FACTOR, MAX_FACTOR, - INTERPOLATOR_POWER, N_STAGES_EXTENDED, ERROR_ESTIMATOR_ORDER, ERROR_EXPONENT, @@ -195,7 +194,7 @@ class Dop853DenseOutput: def __init__(self, t_old, h, rr_old, vv_old, F): self.t_old = t_old self.h = h - self.F = F + self._F = F self.rr_old = rr_old self.vv_old = vv_old @@ -214,13 +213,15 @@ def __call__(self, t: float): 1-D array. """ + F00, F01, F02, F03, F04, F05, F06 = self._F + x = (t - self.t_old) / self.h rr_new = (0.0, 0.0, 0.0) vv_new = (0.0, 0.0, 0.0) - for idx, f in enumerate(reversed(self.F)): - rr_new = add_VV_hf(rr_new, array_to_V_hf(f[:3])) - vv_new = add_VV_hf(vv_new, array_to_V_hf(f[3:])) + for idx, f in enumerate((F06, F05, F04, F03, F02, F01, F00)): + rr_new = add_VV_hf(rr_new, f[:3]) + vv_new = add_VV_hf(vv_new, f[3:]) if idx % 2 == 0: rr_new = mul_Vs_hf(rr_new, x) @@ -388,33 +389,28 @@ def dense_output(self): ) # TODO call into hf K[s] = np.array([*rr, *vv]) - F = np.empty((INTERPOLATOR_POWER, N_RV), dtype=float) # TODO use correct type - fr_old = array_to_V_hf(K[0, :3]) fv_old = array_to_V_hf(K[0, 3:]) delta_rr = sub_VV_hf(self.rr, self.rr_old) delta_vv = sub_VV_hf(self.vv, self.vv_old) - F[0, :3] = delta_rr - F[0, 3:] = delta_vv - - F[1, :3] = sub_VV_hf(mul_Vs_hf(fr_old, h), delta_rr) - F[1, 3:] = sub_VV_hf(mul_Vs_hf(fv_old, h), delta_vv) - - F[2, :3] = sub_VV_hf( - mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(self.fr, fr_old), h) - ) - F[2, 3:] = sub_VV_hf( - mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(self.fv, fv_old), h) + F00 = *delta_rr, *delta_vv + F01 = *sub_VV_hf(mul_Vs_hf(fr_old, h), delta_rr), *sub_VV_hf( + mul_Vs_hf(fv_old, h), delta_vv ) + F02 = *sub_VV_hf( + mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(self.fr, fr_old), h) + ), *sub_VV_hf(mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(self.fv, fv_old), h)) - F[3:, :] = h * np.dot(self.D, K) # TODO + F03, F04, F05, F06 = tuple( + tuple(float(number) for number in line) for line in (h * np.dot(self.D, K)) + ) # TODO return Dop853DenseOutput( self.t_old, self.t - self.t_old, # h self.rr_old, self.vv_old, - F, + (F00, F01, F02, F03, F04, F05, F06), ) From 5afe560f542bd047d61708278e09959d32bf5662 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 08:33:03 +0100 Subject: [PATCH 219/346] isolate ieee754 stuff --- src/hapsira/core/math/ieee754.py | 30 +++++++++++++++++++++++++++ src/hapsira/core/math/ivp/_brentq.py | 2 +- src/hapsira/core/math/ivp/_rkerror.py | 11 +--------- src/hapsira/core/math/ivp/_rkstep.py | 12 ++--------- src/hapsira/core/math/ivp/_solve.py | 2 +- src/hapsira/core/math/linalg.py | 17 +-------------- 6 files changed, 36 insertions(+), 38 deletions(-) create mode 100644 src/hapsira/core/math/ieee754.py diff --git a/src/hapsira/core/math/ieee754.py b/src/hapsira/core/math/ieee754.py new file mode 100644 index 000000000..677925f99 --- /dev/null +++ b/src/hapsira/core/math/ieee754.py @@ -0,0 +1,30 @@ +from numpy import ( + float64 as f8, + float32 as f4, + float16 as f2, + finfo, +) + +from ...settings import settings + + +__all__ = [ + "EPS", + "f8", + "f4", + "f2", + "float_", +] + + +if settings["PRECISION"].value == "f8": + float_ = f8 +elif settings["PRECISION"].value == "f4": + float_ = f4 +elif settings["PRECISION"].value == "f2": + float_ = f2 +else: + raise ValueError("unsupported precision") + + +EPS = finfo(float_).eps diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index af85d891a..0885dcdc5 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -1,6 +1,6 @@ from math import fabs, isnan -from ..linalg import EPS +from ..ieee754 import EPS from ...jit import hjit diff --git a/src/hapsira/core/math/ivp/_rkerror.py b/src/hapsira/core/math/ivp/_rkerror.py index a4eb1d152..5f3900156 100644 --- a/src/hapsira/core/math/ivp/_rkerror.py +++ b/src/hapsira/core/math/ivp/_rkerror.py @@ -2,17 +2,8 @@ from ._const import N_RV, KSIG from ._dop853_coefficients import E3 as _E3, E5 as _E5 +from ..ieee754 import float_ from ...jit import hjit -from ....settings import settings - -if settings["PRECISION"].value == "f8": - from numpy import float64 as float_ -elif settings["PRECISION"].value == "f4": - from numpy import float32 as float_ -elif settings["PRECISION"].value == "f2": - from numpy import float16 as float_ -else: - raise ValueError("unsupported precision") __all__ = [ diff --git a/src/hapsira/core/math/ivp/_rkstep.py b/src/hapsira/core/math/ivp/_rkstep.py index 83d031811..37fbb832d 100644 --- a/src/hapsira/core/math/ivp/_rkstep.py +++ b/src/hapsira/core/math/ivp/_rkstep.py @@ -1,23 +1,15 @@ from ._const import N_STAGES, KSIG from ._dop853_coefficients import A as _A, B as _B, C as _C +from ..ieee754 import float_ from ..linalg import add_VV_hf from ...jit import hjit, DSIG -from ....settings import settings - -if settings["PRECISION"].value == "f8": - from numpy import float64 as float_ -elif settings["PRECISION"].value == "f4": - from numpy import float32 as float_ -elif settings["PRECISION"].value == "f2": - from numpy import float16 as float_ -else: - raise ValueError("unsupported precision") __all__ = [ "rk_step_hf", ] + A01 = tuple(float_(number) for number in _A[1, :N_STAGES]) A02 = tuple(float_(number) for number in _A[2, :N_STAGES]) A03 = tuple(float_(number) for number in _A[3, :N_STAGES]) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 59aa6f76a..9557ef2fc 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -5,7 +5,7 @@ from ._brentq import brentq_hf, BRENTQ_CONVERGED, BRENTQ_MAXITER from ._solution import OdeSolution from ._rk import DOP853 -from ...math.linalg import EPS +from ..ieee754 import EPS __all__ = [ diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 58f962862..226d069f0 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -1,8 +1,7 @@ from math import inf, sqrt -from numpy import finfo +from .ieee754 import EPS from ..jit import hjit, vjit -from ...settings import settings __all__ = [ "abs_V_hf", @@ -24,23 +23,9 @@ "sign_hf", "sub_VV_hf", "transpose_M_hf", - "EPS", ] -if settings["PRECISION"].value == "f8": - from numpy import float64 as float_ -elif settings["PRECISION"].value == "f4": - from numpy import float32 as float_ -elif settings["PRECISION"].value == "f2": - from numpy import float16 as float_ -else: - raise ValueError("unsupported precision") - - -EPS = finfo(float_).eps - - @hjit("V(V)") def abs_V_hf(x): return abs(x[0]), abs(x[1]), abs(x[2]) From e94b8cdcd01a7fbe74742c5bef5977d2ebab00d3 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 08:46:18 +0100 Subject: [PATCH 220/346] structure rk --- src/hapsira/core/math/ivp/_rk.py | 168 +---------------------- src/hapsira/core/math/ivp/_rkstepimpl.py | 107 +++++++++++++++ src/hapsira/core/math/ivp/_rkstepinit.py | 57 ++++++++ 3 files changed, 170 insertions(+), 162 deletions(-) create mode 100644 src/hapsira/core/math/ivp/_rkstepimpl.py create mode 100644 src/hapsira/core/math/ivp/_rkstepinit.py diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index a8e3430b1..042ec01ce 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -1,4 +1,3 @@ -from math import inf, sqrt from typing import Callable import numpy as np @@ -6,28 +5,18 @@ from ._const import ( N_RV, N_STAGES, - KSIG, - SAFETY, - MIN_FACTOR, - MAX_FACTOR, N_STAGES_EXTENDED, ERROR_ESTIMATOR_ORDER, - ERROR_EXPONENT, ) from ._dop853_coefficients import A as _A, C as _C, D as _D -from ._rkstep import rk_step_hf -from ._rkerror import estimate_error_norm_V_hf +from ._rkstepinit import select_initial_step_hf +from ._rkstepimpl import step_impl_hf -from ...jit import array_to_V_hf, hjit, DSIG + +from ...jit import array_to_V_hf from ...math.linalg import ( - abs_V_hf, - add_Vs_hf, add_VV_hf, - div_VV_hf, - max_VV_hf, mul_Vs_hf, - nextafter_hf, - norm_VV_hf, sub_VV_hf, EPS, ) @@ -37,151 +26,6 @@ ] -@hjit(f"f(F({DSIG:s}),f,V,V,f,V,V,f,f,f,f)") -def _select_initial_step_hf( - fun, t0, rr, vv, argk, fr, fv, direction, order, rtol, atol -): - scale_r = ( - atol + abs(rr[0]) * rtol, - atol + abs(rr[1]) * rtol, - atol + abs(rr[2]) * rtol, - ) - scale_v = ( - atol + abs(vv[0]) * rtol, - atol + abs(vv[1]) * rtol, - atol + abs(vv[2]) * rtol, - ) - - factor = 1 / sqrt(6) - d0 = norm_VV_hf(div_VV_hf(rr, scale_r), div_VV_hf(vv, scale_v)) * factor - d1 = norm_VV_hf(div_VV_hf(fr, scale_r), div_VV_hf(fv, scale_v)) * factor - - if d0 < 1e-5 or d1 < 1e-5: - h0 = 1e-6 - else: - h0 = 0.01 * d0 / d1 - - yr1 = add_VV_hf(rr, mul_Vs_hf(fr, h0 * direction)) - yv1 = add_VV_hf(vv, mul_Vs_hf(fv, h0 * direction)) - - fr1, fv1 = fun( - t0 + h0 * direction, - yr1, - yv1, - argk, - ) - - d2 = ( - norm_VV_hf( - div_VV_hf(sub_VV_hf(fr1, fr), scale_r), - div_VV_hf(sub_VV_hf(fv1, fv), scale_v), - ) - / h0 - ) - - if d1 <= 1e-15 and d2 <= 1e-15: - h1 = max(1e-6, h0 * 1e-3) - else: - h1 = (0.01 / max(d1, d2)) ** (1 / (order + 1)) - - return min(100 * h0, h1) - - -@hjit( - f"Tuple([b1,f,f,V,V,f,V,V,{KSIG:s}])" - f"(F({DSIG:s}),f,f,V,V,V,V,f,f,f,f,f,{KSIG:s})" -) -def _step_impl_hf( - fun, argk, t, rr, vv, fr, fv, rtol, atol, direction, h_abs, t_bound, K -): - min_step = 10 * abs(nextafter_hf(t, direction * inf) - t) - - if h_abs < min_step: - h_abs = min_step - - step_accepted = False - step_rejected = False - - while not step_accepted: - if h_abs < min_step: - return ( - False, - 0.0, - 0.0, - (0.0, 0.0, 0.0), - (0.0, 0.0, 0.0), - 0.0, - (0.0, 0.0, 0.0), - (0.0, 0.0, 0.0), - K, - ) - - h = h_abs * direction - t_new = t + h - - if direction * (t_new - t_bound) > 0: - t_new = t_bound - - h = t_new - t - h_abs = abs(h) - - rr_new, vv_new, fr_new, fv_new, K_new = rk_step_hf( - fun, - t, - rr, - vv, - fr, - fv, - h, - argk, - ) - - scale_r = add_Vs_hf( - mul_Vs_hf( - max_VV_hf( - abs_V_hf(rr), - abs_V_hf(rr_new), - ), - rtol, - ), - atol, - ) - scale_v = add_Vs_hf( - mul_Vs_hf( - max_VV_hf( - abs_V_hf(vv), - abs_V_hf(vv_new), - ), - rtol, - ), - atol, - ) - error_norm = estimate_error_norm_V_hf( - K_new, - h, - scale_r, - scale_v, - ) - - if error_norm < 1: - if error_norm == 0: - factor = MAX_FACTOR - else: - factor = min(MAX_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) - - if step_rejected: - factor = min(1, factor) - - h_abs *= factor - - step_accepted = True - else: - h_abs *= max(MIN_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) - step_rejected = True - - return True, h, t_new, rr_new, vv_new, h_abs, fr_new, fv_new, K_new - - class Dop853DenseOutput: """local interpolant over step made by an ODE solver. @@ -290,7 +134,7 @@ def __init__( self.argk, ) # TODO call into hf - self.h_abs = _select_initial_step_hf( + self.h_abs = select_initial_step_hf( self.fun, self.t, self.rr, @@ -325,7 +169,7 @@ def step(self): return t = self.t - success, *rets = _step_impl_hf( + success, *rets = step_impl_hf( self.fun, self.argk, self.t, diff --git a/src/hapsira/core/math/ivp/_rkstepimpl.py b/src/hapsira/core/math/ivp/_rkstepimpl.py new file mode 100644 index 000000000..f39cb86cc --- /dev/null +++ b/src/hapsira/core/math/ivp/_rkstepimpl.py @@ -0,0 +1,107 @@ +from math import inf + +from ._const import ERROR_EXPONENT, KSIG, MAX_FACTOR, MIN_FACTOR, SAFETY +from ._rkstep import rk_step_hf +from ._rkerror import estimate_error_norm_V_hf +from ..linalg import abs_V_hf, add_Vs_hf, max_VV_hf, mul_Vs_hf, nextafter_hf +from ...jit import hjit, DSIG + + +__all__ = [ + "step_impl_hf", +] + + +@hjit( + f"Tuple([b1,f,f,V,V,f,V,V,{KSIG:s}])" + f"(F({DSIG:s}),f,f,V,V,V,V,f,f,f,f,f,{KSIG:s})" +) +def step_impl_hf( + fun, argk, t, rr, vv, fr, fv, rtol, atol, direction, h_abs, t_bound, K +): + min_step = 10 * abs(nextafter_hf(t, direction * inf) - t) + + if h_abs < min_step: + h_abs = min_step + + step_accepted = False + step_rejected = False + + while not step_accepted: + if h_abs < min_step: + return ( + False, + 0.0, + 0.0, + (0.0, 0.0, 0.0), + (0.0, 0.0, 0.0), + 0.0, + (0.0, 0.0, 0.0), + (0.0, 0.0, 0.0), + K, + ) + + h = h_abs * direction + t_new = t + h + + if direction * (t_new - t_bound) > 0: + t_new = t_bound + + h = t_new - t + h_abs = abs(h) + + rr_new, vv_new, fr_new, fv_new, K_new = rk_step_hf( + fun, + t, + rr, + vv, + fr, + fv, + h, + argk, + ) + + scale_r = add_Vs_hf( + mul_Vs_hf( + max_VV_hf( + abs_V_hf(rr), + abs_V_hf(rr_new), + ), + rtol, + ), + atol, + ) + scale_v = add_Vs_hf( + mul_Vs_hf( + max_VV_hf( + abs_V_hf(vv), + abs_V_hf(vv_new), + ), + rtol, + ), + atol, + ) + error_norm = estimate_error_norm_V_hf( + K_new, + h, + scale_r, + scale_v, + ) + + if error_norm < 1: + if error_norm == 0: + factor = MAX_FACTOR + else: + factor = min(MAX_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) + + if step_rejected: + factor = min(1, factor) + + h_abs *= factor + + step_accepted = True + else: + h_abs *= max(MIN_FACTOR, SAFETY * error_norm**ERROR_EXPONENT) + step_rejected = True + + return True, h, t_new, rr_new, vv_new, h_abs, fr_new, fv_new, K_new diff --git a/src/hapsira/core/math/ivp/_rkstepinit.py b/src/hapsira/core/math/ivp/_rkstepinit.py new file mode 100644 index 000000000..b975bec50 --- /dev/null +++ b/src/hapsira/core/math/ivp/_rkstepinit.py @@ -0,0 +1,57 @@ +from math import sqrt + +from ...jit import hjit, DSIG +from ..linalg import add_VV_hf, div_VV_hf, mul_Vs_hf, norm_VV_hf, sub_VV_hf + + +__all__ = [ + "select_initial_step_hf", +] + + +@hjit(f"f(F({DSIG:s}),f,V,V,f,V,V,f,f,f,f)") +def select_initial_step_hf(fun, t0, rr, vv, argk, fr, fv, direction, order, rtol, atol): + scale_r = ( + atol + abs(rr[0]) * rtol, + atol + abs(rr[1]) * rtol, + atol + abs(rr[2]) * rtol, + ) + scale_v = ( + atol + abs(vv[0]) * rtol, + atol + abs(vv[1]) * rtol, + atol + abs(vv[2]) * rtol, + ) + + factor = 1 / sqrt(6) + d0 = norm_VV_hf(div_VV_hf(rr, scale_r), div_VV_hf(vv, scale_v)) * factor + d1 = norm_VV_hf(div_VV_hf(fr, scale_r), div_VV_hf(fv, scale_v)) * factor + + if d0 < 1e-5 or d1 < 1e-5: + h0 = 1e-6 + else: + h0 = 0.01 * d0 / d1 + + yr1 = add_VV_hf(rr, mul_Vs_hf(fr, h0 * direction)) + yv1 = add_VV_hf(vv, mul_Vs_hf(fv, h0 * direction)) + + fr1, fv1 = fun( + t0 + h0 * direction, + yr1, + yv1, + argk, + ) + + d2 = ( + norm_VV_hf( + div_VV_hf(sub_VV_hf(fr1, fr), scale_r), + div_VV_hf(sub_VV_hf(fv1, fv), scale_v), + ) + / h0 + ) + + if d1 <= 1e-15 and d2 <= 1e-15: + h1 = max(1e-6, h0 * 1e-3) + else: + h1 = (0.01 / max(d1, d2)) ** (1 / (order + 1)) + + return min(100 * h0, h1) From bb87d6bb8713d4e81de8dad4626fd7282f83a99c Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 08:52:38 +0100 Subject: [PATCH 221/346] rm K_extended --- src/hapsira/core/math/ivp/_rk.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 042ec01ce..a6c78d9d2 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -116,10 +116,8 @@ def __init__( self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 - self.K_extended = np.empty( - (N_STAGES_EXTENDED, N_RV), dtype=float - ) # TODO set type - self.K = self.K_extended[: N_STAGES + 1, :] + self.K = np.empty((N_STAGES + 1, N_RV), dtype=float) + self.rr_old = None self.vv_old = None self.t_old = None @@ -218,7 +216,9 @@ def dense_output(self): assert self.t_old is not None assert self.t != self.t_old - K = self.K_extended + K = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) + K[: N_STAGES + 1, :] = self.K + h = self.h_previous for s, (a, c) in enumerate(zip(self.A_EXTRA, self.C_EXTRA), start=N_STAGES + 1): From a93c70f691dd7782ee5dfb9be20ac50b09f1398a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 08:54:56 +0100 Subject: [PATCH 222/346] cleanup --- src/hapsira/core/math/ivp/_rk.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index a6c78d9d2..c8324c475 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -180,19 +180,18 @@ def step(self): self.direction, self.h_abs, self.t_bound, - tuple(tuple(line) for line in self.K[: N_STAGES + 1, :N_RV]), + tuple(tuple(line) for line in self.K), ) if success: self.h_previous = rets[0] - # self.y_old = np.array([*rets[1], *rets[2]]) self.rr_old = self.rr self.vv_old = self.vv self.t = rets[1] self.rr, self.vv = rets[2], rets[3] self.h_abs = rets[4] self.fr, self.fv = rets[5], rets[6] - self.K[: N_STAGES + 1, :N_RV] = np.array(rets[7]) + self.K[:, :] = np.array(rets[7]) if not success: self.status = "failed" From 0a123f0013dde8c2605de015865020c77bb3924f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 09:01:42 +0100 Subject: [PATCH 223/346] K becomes tuple --- src/hapsira/core/math/ivp/_rk.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index c8324c475..24b439b0a 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -116,7 +116,21 @@ def __init__( self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 - self.K = np.empty((N_STAGES + 1, N_RV), dtype=float) + self.K = ( + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 0 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 1 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 2 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 3 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 4 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 5 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 6 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 7 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 8 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 9 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 10 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 11 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 12 + ) self.rr_old = None self.vv_old = None @@ -180,7 +194,7 @@ def step(self): self.direction, self.h_abs, self.t_bound, - tuple(tuple(line) for line in self.K), + self.K, ) if success: @@ -191,7 +205,7 @@ def step(self): self.rr, self.vv = rets[2], rets[3] self.h_abs = rets[4] self.fr, self.fv = rets[5], rets[6] - self.K[:, :] = np.array(rets[7]) + self.K = rets[7] if not success: self.status = "failed" @@ -216,7 +230,7 @@ def dense_output(self): assert self.t != self.t_old K = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) - K[: N_STAGES + 1, :] = self.K + K[: N_STAGES + 1, :] = np.array(self.K) h = self.h_previous From 60f5d3e946e3f02c6733359d11d5b976f5abd3ac Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 09:05:05 +0100 Subject: [PATCH 224/346] cleanup --- src/hapsira/core/math/ivp/_rk.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 24b439b0a..5d3eedc0a 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -198,14 +198,18 @@ def step(self): ) if success: - self.h_previous = rets[0] self.rr_old = self.rr self.vv_old = self.vv - self.t = rets[1] - self.rr, self.vv = rets[2], rets[3] - self.h_abs = rets[4] - self.fr, self.fv = rets[5], rets[6] - self.K = rets[7] + ( + self.h_previous, + self.t, + self.rr, + self.vv, + self.h_abs, + self.fr, + self.fv, + self.K, + ) = rets if not success: self.status = "failed" From f575fc33980c135c52461adab862cf404f63294a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 13:46:16 +0100 Subject: [PATCH 225/346] new dense interp --- src/hapsira/core/math/ivp/_rkdenseinterp.py | 63 +++++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 src/hapsira/core/math/ivp/_rkdenseinterp.py diff --git a/src/hapsira/core/math/ivp/_rkdenseinterp.py b/src/hapsira/core/math/ivp/_rkdenseinterp.py new file mode 100644 index 000000000..52ca89bbe --- /dev/null +++ b/src/hapsira/core/math/ivp/_rkdenseinterp.py @@ -0,0 +1,63 @@ +from ._const import FSIG +from ..linalg import add_VV_hf, mul_Vs_hf +from ...jit import hjit + + +__all__ = [ + "dense_interp_brentq_hb", + "dense_interp_hf", + "DENSE_SIG", +] + + +DENSE_SIG = f"f,f,V,V,{FSIG:s}" + + +@hjit(f"Tuple([V,V])(f,{DENSE_SIG:s})") +def dense_interp_hf(t, t_old, h, rr_old, vv_old, F): + """ + Local interpolant over step made by an ODE solver. + Evaluate the interpolant. + + Parameters + ---------- + t : float or array_like with shape (n_points,) + Points to evaluate the solution at. + + Returns + ------- + y : ndarray, shape (n,) or (n, n_points) + Computed values. Shape depends on whether `t` was a scalar or a + 1-D array. + """ + + F00, F01, F02, F03, F04, F05, F06 = F + + x = (t - t_old) / h + rr_new = (0.0, 0.0, 0.0) + vv_new = (0.0, 0.0, 0.0) + + for idx, f in enumerate((F06, F05, F04, F03, F02, F01, F00)): + rr_new = add_VV_hf(rr_new, f[:3]) + vv_new = add_VV_hf(vv_new, f[3:]) + + if idx % 2 == 0: + rr_new = mul_Vs_hf(rr_new, x) + vv_new = mul_Vs_hf(vv_new, x) + else: + rr_new = mul_Vs_hf(rr_new, 1 - x) + vv_new = mul_Vs_hf(vv_new, 1 - x) + + rr_new = add_VV_hf(rr_new, rr_old) + vv_new = add_VV_hf(vv_new, vv_old) + + return rr_new, vv_new + + +def dense_interp_brentq_hb(func): + @hjit(f"f(f,{DENSE_SIG:s},f)", cache=False) + def event_wrapper(t, t_old, h, rr_old, vv_old, F, argk): + rr, vv = dense_interp_hf(t, t_old, h, rr_old, vv_old, F) + return func(t, rr, vv, argk) + + return event_wrapper From 4af915f5516548b1f8490e18756d39d208ae7e12 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 13:46:58 +0100 Subject: [PATCH 226/346] compiled wrapper around events for brentq --- src/hapsira/twobody/events.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index d9bb60e09..59efb68a4 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -6,6 +6,7 @@ from astropy.coordinates import get_body_barycentric_posvel from hapsira.core.jit import hjit +from hapsira.core.math.ivp import dense_interp_brentq_hb from hapsira.core.math.linalg import mul_Vs_hf, norm_V_hf from hapsira.core.events import ( eclipse_function_hf, @@ -46,6 +47,7 @@ def __init__(self, terminal, direction): self._terminal, self._direction = terminal, direction self._last_t = None self._impl_hf = None + self._impl_dense_hf = None @property def terminal(self): @@ -71,10 +73,18 @@ def last_t_raw(self, value: float): def impl_hf(self) -> Callable: return self._impl_hf + @property + def impl_dense_hf(self) -> Callable: + return self._impl_dense_hf + def __call__(self, t, rr, vv, k): + raise NotImplementedError() # HACK self._last_t = t return self._impl_hf(t, rr, vv, k) + def _wrap(self): + self._impl_dense_hf = dense_interp_brentq_hb(self._impl_hf) + class AltitudeCrossEvent(BaseEvent): """Detect if a satellite crosses a specific threshold altitude. @@ -105,6 +115,7 @@ def impl_hf(t, rr, vv, k): ) # If this goes from +ve to -ve, altitude is decreasing. self._impl_hf = impl_hf + self._wrap() class LithobrakeEvent(AltitudeCrossEvent): @@ -154,6 +165,7 @@ def impl_hf(t, rr, vv, k): return rad2deg(lat_) - lat self._impl_hf = impl_hf + self._wrap() class BaseEclipseEvent(BaseEvent): @@ -231,6 +243,7 @@ def impl_hf(t, rr, vv, k): return shadow_function self._impl_hf = impl_hf + self._wrap() class UmbraEvent(BaseEclipseEvent): @@ -270,6 +283,7 @@ def impl_hf(t, rr, vv, k): return shadow_function self._impl_hf = impl_hf + self._wrap() class NodeCrossEvent(BaseEvent): @@ -295,6 +309,7 @@ def impl_hf(t, rr, vv, k): return rr[2] self._impl_hf = impl_hf + self._wrap() class LosEvent(BaseEvent): @@ -331,3 +346,4 @@ def impl_hf(t, rr, vv, k): return delta_angle self._impl_hf = impl_hf + self._wrap() From 98d7fec9bc4b7e33ec9426979015aff05e42cda1 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 13:47:30 +0100 Subject: [PATCH 227/346] rm dense output class --- src/hapsira/core/math/ivp/_rk.py | 56 +------------------------------- 1 file changed, 1 insertion(+), 55 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 5d3eedc0a..60bed5dc2 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -26,60 +26,6 @@ ] -class Dop853DenseOutput: - """local interpolant over step made by an ODE solver. - - Attributes - ---------- - t_min, t_max : float - Time range of the interpolation. - """ - - def __init__(self, t_old, h, rr_old, vv_old, F): - self.t_old = t_old - self.h = h - self._F = F - self.rr_old = rr_old - self.vv_old = vv_old - - def __call__(self, t: float): - """Evaluate the interpolant. - - Parameters - ---------- - t : float or array_like with shape (n_points,) - Points to evaluate the solution at. - - Returns - ------- - y : ndarray, shape (n,) or (n, n_points) - Computed values. Shape depends on whether `t` was a scalar or a - 1-D array. - """ - - F00, F01, F02, F03, F04, F05, F06 = self._F - - x = (t - self.t_old) / self.h - rr_new = (0.0, 0.0, 0.0) - vv_new = (0.0, 0.0, 0.0) - - for idx, f in enumerate((F06, F05, F04, F03, F02, F01, F00)): - rr_new = add_VV_hf(rr_new, f[:3]) - vv_new = add_VV_hf(vv_new, f[3:]) - - if idx % 2 == 0: - rr_new = mul_Vs_hf(rr_new, x) - vv_new = mul_Vs_hf(vv_new, x) - else: - rr_new = mul_Vs_hf(rr_new, 1 - x) - vv_new = mul_Vs_hf(vv_new, 1 - x) - - rr_new = add_VV_hf(rr_new, self.rr_old) - vv_new = add_VV_hf(vv_new, self.vv_old) - - return rr_new, vv_new - - class DOP853: """ Explicit Runge-Kutta method of order 8. @@ -268,7 +214,7 @@ def dense_output(self): tuple(float(number) for number in line) for line in (h * np.dot(self.D, K)) ) # TODO - return Dop853DenseOutput( + return ( self.t_old, self.t - self.t_old, # h self.rr_old, From 84d5ab4a8adb40996b16b04e2858b1df76d631d2 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 13:49:04 +0100 Subject: [PATCH 228/346] run on compiled wrapper, keep track of last_t manually --- src/hapsira/core/math/ivp/_solve.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 9557ef2fc..77530db12 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -2,7 +2,7 @@ import numpy as np -from ._brentq import brentq_hf, BRENTQ_CONVERGED, BRENTQ_MAXITER +from ._brentq import brentq_dense_hf, BRENTQ_CONVERGED, BRENTQ_MAXITER from ._solution import OdeSolution from ._rk import DOP853 from ..ieee754 import EPS @@ -39,18 +39,17 @@ def _solve_event_equation( Found solution. """ - def wrapper(t): - rr, vv = sol(t) - return event(t, rr, vv, argk) - - value, status = brentq_hf( - wrapper, + last_t, value, status = brentq_dense_hf( + event.impl_dense_hf, t_old, t, 4 * EPS, 4 * EPS, BRENTQ_MAXITER, + *sol, + argk, ) + event.last_t_raw = last_t assert BRENTQ_CONVERGED == status return value @@ -211,7 +210,10 @@ def solve_ivp( events, is_terminal, event_dir = _prepare_events(events) if events is not None: - g = [event(t0, rr, vv, argk) for event in events] + g = [] + for event in events: + g.append(event.impl_hf(t0, rr, vv, argk)) + event.last_t_raw = t0 status = None while status is None: @@ -230,7 +232,10 @@ def solve_ivp( interpolants.append(sol) if events is not None: - g_new = [event(t, solver.rr, solver.vv, argk) for event in events] + g_new = [] + for event in events: + g_new.append(event.impl_hf(t, solver.rr, solver.vv, argk)) + event.last_t_raw = t active_events = _find_active_events(g, g_new, event_dir) if active_events.size > 0: _, roots, terminate = _handle_events( From 1c14927899b3f59f76d9a3c54344766a8ab81131 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 13:50:00 +0100 Subject: [PATCH 229/346] events are not directly callable anymore, compiled implementation to be used instead --- src/hapsira/twobody/events.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index 59efb68a4..6b0783e3b 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -77,11 +77,6 @@ def impl_hf(self) -> Callable: def impl_dense_hf(self) -> Callable: return self._impl_dense_hf - def __call__(self, t, rr, vv, k): - raise NotImplementedError() # HACK - self._last_t = t - return self._impl_hf(t, rr, vv, k) - def _wrap(self): self._impl_dense_hf = dense_interp_brentq_hb(self._impl_hf) From c2ce6e2f8040ea437bf77b2270849329d4197481 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 13:50:41 +0100 Subject: [PATCH 230/346] secondary brentq implementation that wrapps dense output funcs --- src/hapsira/core/math/ivp/_brentq.py | 104 ++++++++++++++++++++++++++- 1 file changed, 103 insertions(+), 1 deletion(-) diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index 0885dcdc5..db6c2846f 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -1,5 +1,6 @@ -from math import fabs, isnan +from math import fabs, isnan, nan +from ._rkdenseinterp import DENSE_SIG from ..ieee754 import EPS from ...jit import hjit @@ -13,6 +14,7 @@ "BRENTQ_RTOL", "BRENTQ_MAXITER", "brentq_hf", + "brentq_dense_hf", ] @@ -128,3 +130,103 @@ def brentq_hf( return 0.0, BRENTQ_ERROR return xcur, BRENTQ_CONVERR + + +@hjit(f"Tuple([f,f,i8])(F(f(f,{DENSE_SIG:s},f)),f,f,f,f,f,{DENSE_SIG:s},f)") +def brentq_dense_hf( + func, # callback_type + xa, # double + xb, # double + xtol, # double + rtol, # double + maxiter, # int + sol1, + sol2, + sol3, + sol4, + sol5, + argk, +): + """ + Loosely adapted from + https://github.com/scipy/scipy/blob/d23363809572e9a44074a3f06f66137083446b48/scipy/optimize/_zeros_py.py#L682 + """ + + if not xtol + 0.0 > 0: + return nan, 0.0, BRENTQ_ERROR + if not rtol + 0.0 >= BRENTQ_RTOL: + return nan, 0.0, BRENTQ_ERROR + if not maxiter + 0 >= 0: + return nan, 0.0, BRENTQ_ERROR + + xpre, xcur = xa, xb + xblk = 0.0 + fpre, fcur, fblk = 0.0, 0.0, 0.0 + spre, scur = 0.0, 0.0 + + fpre = func(xpre, sol1, sol2, sol3, sol4, sol5, argk) + if isnan(fpre): + return xpre, 0.0, BRENTQ_ERROR + + fcur = func(xcur, sol1, sol2, sol3, sol4, sol5, argk) + if isnan(fcur): + return xcur, 0.0, BRENTQ_ERROR + + if fpre == 0: + return xcur, xpre, BRENTQ_CONVERGED + if fcur == 0: + return xcur, xcur, BRENTQ_CONVERGED + if _signbit_s_hf(fpre) == _signbit_s_hf(fcur): + return xcur, 0.0, BRENTQ_SIGNERR + + for _ in range(0, maxiter): + if fpre != 0 and fcur != 0 and _signbit_s_hf(fpre) != _signbit_s_hf(fcur): + xblk = xpre + fblk = fpre + scur = xcur - xpre + spre = scur + if fabs(fblk) < fabs(fcur): + xpre = xcur + xcur = xblk + xblk = xpre + + fpre = fcur + fcur = fblk + fblk = fpre + + delta = (xtol + rtol * fabs(xcur)) / 2 + sbis = (xblk - xcur) / 2 + if fcur == 0 or fabs(sbis) < delta: + return xcur, xcur, BRENTQ_CONVERGED + + if fabs(spre) > delta and fabs(fcur) < fabs(fpre): + if xpre == xblk: + stry = -fcur * (xcur - xpre) / (fcur - fpre) + else: + dpre = (fpre - fcur) / (xpre - xcur) + dblk = (fblk - fcur) / (xblk - xcur) + stry = ( + -fcur * (fblk * dblk - fpre * dpre) / (dblk * dpre * (fblk - fpre)) + ) + if 2 * fabs(stry) < _min_ss_hf(fabs(spre), 3 * fabs(sbis) - delta): + spre = scur + scur = stry + else: + spre = sbis + scur = sbis + else: + spre = sbis + scur = sbis + + xpre = xcur + fpre = fcur + if fabs(scur) > delta: + xcur += scur + else: + xcur += delta if sbis > 0 else -delta + + fcur = func(xcur, sol1, sol2, sol3, sol4, sol5, argk) + if isnan(fcur): + return xcur, 0.0, BRENTQ_ERROR + + return xcur, xcur, BRENTQ_CONVERR From b91442b1fa2c92a8695dc9d1b1bf6ed5080a50c4 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 13:51:16 +0100 Subject: [PATCH 231/346] expose dense output wrappers --- src/hapsira/core/math/ivp/__init__.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/math/ivp/__init__.py b/src/hapsira/core/math/ivp/__init__.py index 138cec664..f918f2b4d 100644 --- a/src/hapsira/core/math/ivp/__init__.py +++ b/src/hapsira/core/math/ivp/__init__.py @@ -1,4 +1,3 @@ -from ._solve import solve_ivp from ._brentq import ( BRENTQ_CONVERGED, BRENTQ_SIGNERR, @@ -9,9 +8,10 @@ BRENTQ_MAXITER, brentq_hf, ) +from ._rkdenseinterp import dense_interp_brentq_hb, dense_interp_hf +from ._solve import solve_ivp __all__ = [ - "solve_ivp", "BRENTQ_CONVERGED", "BRENTQ_SIGNERR", "BRENTQ_CONVERR", @@ -20,4 +20,7 @@ "BRENTQ_RTOL", "BRENTQ_MAXITER", "brentq_hf", + "dense_interp_brentq_hb", + "dense_interp_hf", + "solve_ivp", ] From eae5d54134520783d405db4fef836cc98b9b9f38 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 13:51:37 +0100 Subject: [PATCH 232/346] use new dense output wrapper --- src/hapsira/core/math/ivp/_solution.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solution.py b/src/hapsira/core/math/ivp/_solution.py index a78205d62..0ed38344a 100644 --- a/src/hapsira/core/math/ivp/_solution.py +++ b/src/hapsira/core/math/ivp/_solution.py @@ -1,5 +1,7 @@ import numpy as np +from ._rkdenseinterp import dense_interp_hf + class OdeSolution: """Continuous ODE solution. @@ -65,10 +67,10 @@ def __call__(self, t): Computed values. Shape depends on whether `t` is a scalar or a 1-D array. """ - t = np.asarray(t) + # t = np.asarray(t) - assert t.ndim == 0 + # assert t.ndim == 0 ind = np.searchsorted(self.ts_sorted, t, side="left") segment = min(max(ind - 1, 0), self.n_segments - 1) - return self.interpolants[segment](t) + return dense_interp_hf(t, *self.interpolants[segment]) From aba2cecaf2be06625013a12889670cec7eb4a788 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 13:52:31 +0100 Subject: [PATCH 233/346] dense parameter signature --- src/hapsira/core/math/ivp/_const.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/hapsira/core/math/ivp/_const.py b/src/hapsira/core/math/ivp/_const.py index 7ed0b3bc4..f2bc7ea7d 100644 --- a/src/hapsira/core/math/ivp/_const.py +++ b/src/hapsira/core/math/ivp/_const.py @@ -30,3 +30,9 @@ + ",".join(["Tuple([" + ",".join(["f"] * N_RV) + "])"] * (N_STAGES + 1)) + "])" ) + +FSIG = ( + "Tuple([" + + ",".join(["Tuple([" + ",".join(["f"] * N_RV) + "])"] * INTERPOLATOR_POWER) + + "])" +) From 43d04d2434e27749204c10a04d8a2026a8a22f9c Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 13:59:46 +0100 Subject: [PATCH 234/346] unroll loop --- src/hapsira/core/math/ivp/_rkdenseinterp.py | 46 +++++++++++++++------ 1 file changed, 33 insertions(+), 13 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseinterp.py b/src/hapsira/core/math/ivp/_rkdenseinterp.py index 52ca89bbe..9c6aa442b 100644 --- a/src/hapsira/core/math/ivp/_rkdenseinterp.py +++ b/src/hapsira/core/math/ivp/_rkdenseinterp.py @@ -34,19 +34,39 @@ def dense_interp_hf(t, t_old, h, rr_old, vv_old, F): F00, F01, F02, F03, F04, F05, F06 = F x = (t - t_old) / h - rr_new = (0.0, 0.0, 0.0) - vv_new = (0.0, 0.0, 0.0) - - for idx, f in enumerate((F06, F05, F04, F03, F02, F01, F00)): - rr_new = add_VV_hf(rr_new, f[:3]) - vv_new = add_VV_hf(vv_new, f[3:]) - - if idx % 2 == 0: - rr_new = mul_Vs_hf(rr_new, x) - vv_new = mul_Vs_hf(vv_new, x) - else: - rr_new = mul_Vs_hf(rr_new, 1 - x) - vv_new = mul_Vs_hf(vv_new, 1 - x) + + rr_new = mul_Vs_hf(F06[:3], x) + vv_new = mul_Vs_hf(F06[3:], x) + + rr_new = add_VV_hf(rr_new, F05[:3]) + vv_new = add_VV_hf(vv_new, F05[3:]) + rr_new = mul_Vs_hf(rr_new, 1 - x) + vv_new = mul_Vs_hf(vv_new, 1 - x) + + rr_new = add_VV_hf(rr_new, F04[:3]) + vv_new = add_VV_hf(vv_new, F04[3:]) + rr_new = mul_Vs_hf(rr_new, x) + vv_new = mul_Vs_hf(vv_new, x) + + rr_new = add_VV_hf(rr_new, F03[:3]) + vv_new = add_VV_hf(vv_new, F03[3:]) + rr_new = mul_Vs_hf(rr_new, 1 - x) + vv_new = mul_Vs_hf(vv_new, 1 - x) + + rr_new = add_VV_hf(rr_new, F02[:3]) + vv_new = add_VV_hf(vv_new, F02[3:]) + rr_new = mul_Vs_hf(rr_new, x) + vv_new = mul_Vs_hf(vv_new, x) + + rr_new = add_VV_hf(rr_new, F01[:3]) + vv_new = add_VV_hf(vv_new, F01[3:]) + rr_new = mul_Vs_hf(rr_new, 1 - x) + vv_new = mul_Vs_hf(vv_new, 1 - x) + + rr_new = add_VV_hf(rr_new, F00[:3]) + vv_new = add_VV_hf(vv_new, F00[3:]) + rr_new = mul_Vs_hf(rr_new, x) + vv_new = mul_Vs_hf(vv_new, x) rr_new = add_VV_hf(rr_new, rr_old) vv_new = add_VV_hf(vv_new, vv_old) From 22b22426e0a254d6f016a2926b569426ea959a39 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 14:23:31 +0100 Subject: [PATCH 235/346] isolate dense output prep --- src/hapsira/core/math/ivp/_rk.py | 110 +++++++++++++++------------- src/hapsira/core/math/ivp/_solve.py | 17 ++++- 2 files changed, 73 insertions(+), 54 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index 60bed5dc2..b6641698e 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -23,18 +23,20 @@ __all__ = [ "DOP853", + "dense_output_hf", ] +A_EXTRA = _A[N_STAGES + 1 :] +C_EXTRA = _C[N_STAGES + 1 :] +D = _D + + class DOP853: """ Explicit Runge-Kutta method of order 8. """ - A_EXTRA = _A[N_STAGES + 1 :] - C_EXTRA = _C[N_STAGES + 1 :] - D = _D - def __init__( self, fun: Callable, @@ -167,57 +169,61 @@ def step(self): self.status = "finished" - def dense_output(self): - """Compute a local interpolant over the last successful step. - - Returns - ------- - sol : `DenseOutput` - Local interpolant over the last successful step. - """ - - assert self.t_old is not None - assert self.t != self.t_old - K = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) - K[: N_STAGES + 1, :] = np.array(self.K) +# TODO compile +def dense_output_hf( + fun, argk, t_old, t, h_previous, rr, vv, rr_old, vv_old, fr, fv, K_ +): + """Compute a local interpolant over the last successful step. - h = self.h_previous + Returns + ------- + sol : `DenseOutput` + Local interpolant over the last successful step. + """ - for s, (a, c) in enumerate(zip(self.A_EXTRA, self.C_EXTRA), start=N_STAGES + 1): - dy = np.dot(K[:s].T, a[:s]) * h - rr_ = add_VV_hf(self.rr_old, array_to_V_hf(dy[:3])) - vv_ = add_VV_hf(self.vv_old, array_to_V_hf(dy[3:])) - rr, vv = self.fun( - self.t_old + c * h, - rr_, - vv_, - self.argk, - ) # TODO call into hf - K[s] = np.array([*rr, *vv]) + assert t_old is not None + assert t != t_old - fr_old = array_to_V_hf(K[0, :3]) - fv_old = array_to_V_hf(K[0, 3:]) + Ke = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) + Ke[: N_STAGES + 1, :] = np.array(K_) - delta_rr = sub_VV_hf(self.rr, self.rr_old) - delta_vv = sub_VV_hf(self.vv, self.vv_old) + h = h_previous - F00 = *delta_rr, *delta_vv - F01 = *sub_VV_hf(mul_Vs_hf(fr_old, h), delta_rr), *sub_VV_hf( - mul_Vs_hf(fv_old, h), delta_vv - ) - F02 = *sub_VV_hf( - mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(self.fr, fr_old), h) - ), *sub_VV_hf(mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(self.fv, fv_old), h)) - - F03, F04, F05, F06 = tuple( - tuple(float(number) for number in line) for line in (h * np.dot(self.D, K)) - ) # TODO - - return ( - self.t_old, - self.t - self.t_old, # h - self.rr_old, - self.vv_old, - (F00, F01, F02, F03, F04, F05, F06), - ) + for s, (a, c) in enumerate(zip(A_EXTRA, C_EXTRA), start=N_STAGES + 1): + dy = np.dot(Ke[:s].T, a[:s]) * h + rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) + vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) + rr_, vv_ = fun( + t_old + c * h, + rr_, + vv_, + argk, + ) # TODO call into hf + Ke[s] = np.array([*rr_, *vv_]) + + fr_old = array_to_V_hf(Ke[0, :3]) + fv_old = array_to_V_hf(Ke[0, 3:]) + + delta_rr = sub_VV_hf(rr, rr_old) + delta_vv = sub_VV_hf(vv, vv_old) + + F00 = *delta_rr, *delta_vv + F01 = *sub_VV_hf(mul_Vs_hf(fr_old, h), delta_rr), *sub_VV_hf( + mul_Vs_hf(fv_old, h), delta_vv + ) + F02 = *sub_VV_hf( + mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(fr, fr_old), h) + ), *sub_VV_hf(mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(fv, fv_old), h)) + + F03, F04, F05, F06 = tuple( + tuple(float(number) for number in line) for line in (h * np.dot(D, Ke)) + ) # TODO + + return ( + t_old, + t - t_old, # h + rr_old, + vv_old, + (F00, F01, F02, F03, F04, F05, F06), + ) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 77530db12..46441b54c 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -4,7 +4,7 @@ from ._brentq import brentq_dense_hf, BRENTQ_CONVERGED, BRENTQ_MAXITER from ._solution import OdeSolution -from ._rk import DOP853 +from ._rk import DOP853, dense_output_hf from ..ieee754 import EPS @@ -228,7 +228,20 @@ def solve_ivp( t_old = solver.t_old t = solver.t - sol = solver.dense_output() + sol = dense_output_hf( + solver.fun, + solver.argk, + solver.t_old, + solver.t, + solver.h_previous, + solver.rr, + solver.vv, + solver.rr_old, + solver.vv_old, + solver.fr, + solver.fv, + solver.K, + ) interpolants.append(sol) if events is not None: From 953ea99f058d9bfc81be40aaa891da9d01ce0d70 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 14:29:29 +0100 Subject: [PATCH 236/346] isolate dense output further --- src/hapsira/core/math/ivp/_rk.py | 83 +-------------------- src/hapsira/core/math/ivp/_rkdenseoutput.py | 82 ++++++++++++++++++++ src/hapsira/core/math/ivp/_solve.py | 3 +- 3 files changed, 86 insertions(+), 82 deletions(-) create mode 100644 src/hapsira/core/math/ivp/_rkdenseoutput.py diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rk.py index b6641698e..d8afa0da2 100644 --- a/src/hapsira/core/math/ivp/_rk.py +++ b/src/hapsira/core/math/ivp/_rk.py @@ -2,36 +2,16 @@ import numpy as np -from ._const import ( - N_RV, - N_STAGES, - N_STAGES_EXTENDED, - ERROR_ESTIMATOR_ORDER, -) -from ._dop853_coefficients import A as _A, C as _C, D as _D +from ._const import ERROR_ESTIMATOR_ORDER from ._rkstepinit import select_initial_step_hf from ._rkstepimpl import step_impl_hf - - -from ...jit import array_to_V_hf -from ...math.linalg import ( - add_VV_hf, - mul_Vs_hf, - sub_VV_hf, - EPS, -) +from ..ieee754 import EPS __all__ = [ "DOP853", - "dense_output_hf", ] -A_EXTRA = _A[N_STAGES + 1 :] -C_EXTRA = _C[N_STAGES + 1 :] -D = _D - - class DOP853: """ Explicit Runge-Kutta method of order 8. @@ -168,62 +148,3 @@ def step(self): return self.status = "finished" - - -# TODO compile -def dense_output_hf( - fun, argk, t_old, t, h_previous, rr, vv, rr_old, vv_old, fr, fv, K_ -): - """Compute a local interpolant over the last successful step. - - Returns - ------- - sol : `DenseOutput` - Local interpolant over the last successful step. - """ - - assert t_old is not None - assert t != t_old - - Ke = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) - Ke[: N_STAGES + 1, :] = np.array(K_) - - h = h_previous - - for s, (a, c) in enumerate(zip(A_EXTRA, C_EXTRA), start=N_STAGES + 1): - dy = np.dot(Ke[:s].T, a[:s]) * h - rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) - vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) - rr_, vv_ = fun( - t_old + c * h, - rr_, - vv_, - argk, - ) # TODO call into hf - Ke[s] = np.array([*rr_, *vv_]) - - fr_old = array_to_V_hf(Ke[0, :3]) - fv_old = array_to_V_hf(Ke[0, 3:]) - - delta_rr = sub_VV_hf(rr, rr_old) - delta_vv = sub_VV_hf(vv, vv_old) - - F00 = *delta_rr, *delta_vv - F01 = *sub_VV_hf(mul_Vs_hf(fr_old, h), delta_rr), *sub_VV_hf( - mul_Vs_hf(fv_old, h), delta_vv - ) - F02 = *sub_VV_hf( - mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(fr, fr_old), h) - ), *sub_VV_hf(mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(fv, fv_old), h)) - - F03, F04, F05, F06 = tuple( - tuple(float(number) for number in line) for line in (h * np.dot(D, Ke)) - ) # TODO - - return ( - t_old, - t - t_old, # h - rr_old, - vv_old, - (F00, F01, F02, F03, F04, F05, F06), - ) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py new file mode 100644 index 000000000..e690c51f7 --- /dev/null +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -0,0 +1,82 @@ +import numpy as np + +from ._const import ( + N_RV, + N_STAGES, + N_STAGES_EXTENDED, +) +from ._dop853_coefficients import A as _A, C as _C, D as _D +from ...jit import array_to_V_hf +from ...math.linalg import ( + add_VV_hf, + mul_Vs_hf, + sub_VV_hf, +) + +__all__ = [ + "dense_output_hf", +] + + +A_EXTRA = _A[N_STAGES + 1 :] +C_EXTRA = _C[N_STAGES + 1 :] +D = _D + + +# TODO compile +def dense_output_hf( + fun, argk, t_old, t, h_previous, rr, vv, rr_old, vv_old, fr, fv, K_ +): + """Compute a local interpolant over the last successful step. + + Returns + ------- + sol : `DenseOutput` + Local interpolant over the last successful step. + """ + + assert t_old is not None + assert t != t_old + + Ke = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) + Ke[: N_STAGES + 1, :] = np.array(K_) + + h = h_previous + + for s, (a, c) in enumerate(zip(A_EXTRA, C_EXTRA), start=N_STAGES + 1): + dy = np.dot(Ke[:s].T, a[:s]) * h + rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) + vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) + rr_, vv_ = fun( + t_old + c * h, + rr_, + vv_, + argk, + ) # TODO call into hf + Ke[s] = np.array([*rr_, *vv_]) + + fr_old = array_to_V_hf(Ke[0, :3]) + fv_old = array_to_V_hf(Ke[0, 3:]) + + delta_rr = sub_VV_hf(rr, rr_old) + delta_vv = sub_VV_hf(vv, vv_old) + + F00 = *delta_rr, *delta_vv + F01 = *sub_VV_hf(mul_Vs_hf(fr_old, h), delta_rr), *sub_VV_hf( + mul_Vs_hf(fv_old, h), delta_vv + ) + F02 = *sub_VV_hf( + mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(fr, fr_old), h) + ), *sub_VV_hf(mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(fv, fv_old), h)) + + F03, F04, F05, F06 = tuple( + tuple(float(number) for number in line) for line in (h * np.dot(D, Ke)) + ) # TODO + + return ( + t_old, + t - t_old, # h + rr_old, + vv_old, + (F00, F01, F02, F03, F04, F05, F06), + ) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 46441b54c..01f3d2b62 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -4,7 +4,8 @@ from ._brentq import brentq_dense_hf, BRENTQ_CONVERGED, BRENTQ_MAXITER from ._solution import OdeSolution -from ._rk import DOP853, dense_output_hf +from ._rk import DOP853 +from ._rkdenseoutput import dense_output_hf from ..ieee754 import EPS From 7289753c3b00026afdc121a4dbb7da63047e6821 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 14:36:24 +0100 Subject: [PATCH 237/346] cleanup --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index e690c51f7..28582c3d6 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -6,12 +6,12 @@ N_STAGES_EXTENDED, ) from ._dop853_coefficients import A as _A, C as _C, D as _D -from ...jit import array_to_V_hf -from ...math.linalg import ( +from ..linalg import ( add_VV_hf, mul_Vs_hf, sub_VV_hf, ) +from ...jit import array_to_V_hf __all__ = [ "dense_output_hf", @@ -24,9 +24,7 @@ # TODO compile -def dense_output_hf( - fun, argk, t_old, t, h_previous, rr, vv, rr_old, vv_old, fr, fv, K_ -): +def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): """Compute a local interpolant over the last successful step. Returns @@ -41,8 +39,6 @@ def dense_output_hf( Ke = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) Ke[: N_STAGES + 1, :] = np.array(K_) - h = h_previous - for s, (a, c) in enumerate(zip(A_EXTRA, C_EXTRA), start=N_STAGES + 1): dy = np.dot(Ke[:s].T, a[:s]) * h rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) From e62f66b561a6d18e86b8d368440f3ad8690d3675 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 14:50:13 +0100 Subject: [PATCH 238/346] loop unroll --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 43 +++++++++++++++------ 1 file changed, 32 insertions(+), 11 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index 28582c3d6..9f2dd9c08 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -39,17 +39,38 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): Ke = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) Ke[: N_STAGES + 1, :] = np.array(K_) - for s, (a, c) in enumerate(zip(A_EXTRA, C_EXTRA), start=N_STAGES + 1): - dy = np.dot(Ke[:s].T, a[:s]) * h - rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) - vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) - rr_, vv_ = fun( - t_old + c * h, - rr_, - vv_, - argk, - ) # TODO call into hf - Ke[s] = np.array([*rr_, *vv_]) + dy = np.dot(Ke[:13].T, A_EXTRA[0, :13]) * h + rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) + vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) + rr_, vv_ = fun( + t_old + C_EXTRA[0] * h, + rr_, + vv_, + argk, + ) # TODO call into hf + Ke[13] = np.array([*rr_, *vv_]) + + dy = np.dot(Ke[:14].T, A_EXTRA[1, :14]) * h + rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) + vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) + rr_, vv_ = fun( + t_old + C_EXTRA[1] * h, + rr_, + vv_, + argk, + ) # TODO call into hf + Ke[14] = np.array([*rr_, *vv_]) + + dy = np.dot(Ke[:15].T, A_EXTRA[2, :15]) * h + rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) + vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) + rr_, vv_ = fun( + t_old + C_EXTRA[2] * h, + rr_, + vv_, + argk, + ) # TODO call into hf + Ke[15] = np.array([*rr_, *vv_]) fr_old = array_to_V_hf(Ke[0, :3]) fv_old = array_to_V_hf(Ke[0, 3:]) From ea2ace58ea52cd4d1cc55a77b9a69e7f104a8fa6 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 14:51:12 +0100 Subject: [PATCH 239/346] cleanup --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index 9f2dd9c08..d68075aa7 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -47,7 +47,7 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): rr_, vv_, argk, - ) # TODO call into hf + ) Ke[13] = np.array([*rr_, *vv_]) dy = np.dot(Ke[:14].T, A_EXTRA[1, :14]) * h @@ -58,7 +58,7 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): rr_, vv_, argk, - ) # TODO call into hf + ) Ke[14] = np.array([*rr_, *vv_]) dy = np.dot(Ke[:15].T, A_EXTRA[2, :15]) * h @@ -69,7 +69,7 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): rr_, vv_, argk, - ) # TODO call into hf + ) Ke[15] = np.array([*rr_, *vv_]) fr_old = array_to_V_hf(Ke[0, :3]) From 44c54f4b7623d9b662b9e8baa6477ba8f386bde5 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 14:56:53 +0100 Subject: [PATCH 240/346] arrays to tuples --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index d68075aa7..528cf4159 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -6,6 +6,7 @@ N_STAGES_EXTENDED, ) from ._dop853_coefficients import A as _A, C as _C, D as _D +from ..ieee754 import float_ from ..linalg import ( add_VV_hf, mul_Vs_hf, @@ -18,7 +19,9 @@ ] -A_EXTRA = _A[N_STAGES + 1 :] +A00 = tuple(float_(number) for number in _A[N_STAGES + 1, :13]) +A01 = tuple(float_(number) for number in _A[N_STAGES + 2, :14]) +A02 = tuple(float_(number) for number in _A[N_STAGES + 3, :15]) C_EXTRA = _C[N_STAGES + 1 :] D = _D @@ -39,7 +42,7 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): Ke = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) Ke[: N_STAGES + 1, :] = np.array(K_) - dy = np.dot(Ke[:13].T, A_EXTRA[0, :13]) * h + dy = np.dot(Ke[:13].T, np.array(A00)) * h rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) rr_, vv_ = fun( @@ -50,7 +53,7 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): ) Ke[13] = np.array([*rr_, *vv_]) - dy = np.dot(Ke[:14].T, A_EXTRA[1, :14]) * h + dy = np.dot(Ke[:14].T, np.array(A01)) * h rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) rr_, vv_ = fun( @@ -61,7 +64,7 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): ) Ke[14] = np.array([*rr_, *vv_]) - dy = np.dot(Ke[:15].T, A_EXTRA[2, :15]) * h + dy = np.dot(Ke[:15].T, np.array(A02)) * h rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) rr_, vv_ = fun( From 377c5079d697ec6c9d97226186ea704059561085 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 14:58:50 +0100 Subject: [PATCH 241/346] arrays to tuples --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index 528cf4159..9c34a43a2 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -22,7 +22,7 @@ A00 = tuple(float_(number) for number in _A[N_STAGES + 1, :13]) A01 = tuple(float_(number) for number in _A[N_STAGES + 2, :14]) A02 = tuple(float_(number) for number in _A[N_STAGES + 3, :15]) -C_EXTRA = _C[N_STAGES + 1 :] +C_EXTRA = tuple(float_(number) for number in _C[N_STAGES + 1 :]) D = _D From c7b228659dd301e6f258ecc2e3e6c0658ca879e4 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 15:11:10 +0100 Subject: [PATCH 242/346] replace np.dot --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 492 +++++++++++++++++++- 1 file changed, 488 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index 9c34a43a2..bb1662e02 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -23,7 +23,10 @@ A01 = tuple(float_(number) for number in _A[N_STAGES + 2, :14]) A02 = tuple(float_(number) for number in _A[N_STAGES + 3, :15]) C_EXTRA = tuple(float_(number) for number in _C[N_STAGES + 1 :]) -D = _D +D00 = tuple(float_(number) for number in _D[0, :]) +D01 = tuple(float_(number) for number in _D[1, :]) +D02 = tuple(float_(number) for number in _D[2, :]) +D03 = tuple(float_(number) for number in _D[3, :]) # TODO compile @@ -89,9 +92,490 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(fr, fr_old), h) ), *sub_VV_hf(mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(fv, fv_old), h)) - F03, F04, F05, F06 = tuple( - tuple(float(number) for number in line) for line in (h * np.dot(D, Ke)) - ) # TODO + K00 = tuple(float_(number) for number in Ke[0, :]) + K01 = tuple(float_(number) for number in Ke[1, :]) + K02 = tuple(float_(number) for number in Ke[2, :]) + K03 = tuple(float_(number) for number in Ke[3, :]) + K04 = tuple(float_(number) for number in Ke[4, :]) + K05 = tuple(float_(number) for number in Ke[5, :]) + K06 = tuple(float_(number) for number in Ke[6, :]) + K07 = tuple(float_(number) for number in Ke[7, :]) + K08 = tuple(float_(number) for number in Ke[8, :]) + K09 = tuple(float_(number) for number in Ke[9, :]) + K10 = tuple(float_(number) for number in Ke[10, :]) + K11 = tuple(float_(number) for number in Ke[11, :]) + K12 = tuple(float_(number) for number in Ke[12, :]) + K13 = tuple(float_(number) for number in Ke[13, :]) + K14 = tuple(float_(number) for number in Ke[14, :]) + K15 = tuple(float_(number) for number in Ke[15, :]) + + F03 = ( + ( + D00[0] * K00[0] + + D00[1] * K01[0] + + D00[2] * K02[0] + + D00[3] * K03[0] + + D00[4] * K04[0] + + D00[5] * K05[0] + + D00[6] * K06[0] + + D00[7] * K07[0] + + D00[8] * K08[0] + + D00[9] * K09[0] + + D00[10] * K10[0] + + D00[11] * K11[0] + + D00[12] * K12[0] + + D00[13] * K13[0] + + D00[14] * K14[0] + + D00[15] * K15[0] + ) + * h, + ( + D00[0] * K00[1] + + D00[1] * K01[1] + + D00[2] * K02[1] + + D00[3] * K03[1] + + D00[4] * K04[1] + + D00[5] * K05[1] + + D00[6] * K06[1] + + D00[7] * K07[1] + + D00[8] * K08[1] + + D00[9] * K09[1] + + D00[10] * K10[1] + + D00[11] * K11[1] + + D00[12] * K12[1] + + D00[13] * K13[1] + + D00[14] * K14[1] + + D00[15] * K15[1] + ) + * h, + ( + D00[0] * K00[2] + + D00[1] * K01[2] + + D00[2] * K02[2] + + D00[3] * K03[2] + + D00[4] * K04[2] + + D00[5] * K05[2] + + D00[6] * K06[2] + + D00[7] * K07[2] + + D00[8] * K08[2] + + D00[9] * K09[2] + + D00[10] * K10[2] + + D00[11] * K11[2] + + D00[12] * K12[2] + + D00[13] * K13[2] + + D00[14] * K14[2] + + D00[15] * K15[2] + ) + * h, + ( + D00[0] * K00[3] + + D00[1] * K01[3] + + D00[2] * K02[3] + + D00[3] * K03[3] + + D00[4] * K04[3] + + D00[5] * K05[3] + + D00[6] * K06[3] + + D00[7] * K07[3] + + D00[8] * K08[3] + + D00[9] * K09[3] + + D00[10] * K10[3] + + D00[11] * K11[3] + + D00[12] * K12[3] + + D00[13] * K13[3] + + D00[14] * K14[3] + + D00[15] * K15[3] + ) + * h, + ( + D00[0] * K00[4] + + D00[1] * K01[4] + + D00[2] * K02[4] + + D00[3] * K03[4] + + D00[4] * K04[4] + + D00[5] * K05[4] + + D00[6] * K06[4] + + D00[7] * K07[4] + + D00[8] * K08[4] + + D00[9] * K09[4] + + D00[10] * K10[4] + + D00[11] * K11[4] + + D00[12] * K12[4] + + D00[13] * K13[4] + + D00[14] * K14[4] + + D00[15] * K15[4] + ) + * h, + ( + D00[0] * K00[5] + + D00[1] * K01[5] + + D00[2] * K02[5] + + D00[3] * K03[5] + + D00[4] * K04[5] + + D00[5] * K05[5] + + D00[6] * K06[5] + + D00[7] * K07[5] + + D00[8] * K08[5] + + D00[9] * K09[5] + + D00[10] * K10[5] + + D00[11] * K11[5] + + D00[12] * K12[5] + + D00[13] * K13[5] + + D00[14] * K14[5] + + D00[15] * K15[5] + ) + * h, + ) + + F04 = ( + ( + D01[0] * K00[0] + + D01[1] * K01[0] + + D01[2] * K02[0] + + D01[3] * K03[0] + + D01[4] * K04[0] + + D01[5] * K05[0] + + D01[6] * K06[0] + + D01[7] * K07[0] + + D01[8] * K08[0] + + D01[9] * K09[0] + + D01[10] * K10[0] + + D01[11] * K11[0] + + D01[12] * K12[0] + + D01[13] * K13[0] + + D01[14] * K14[0] + + D01[15] * K15[0] + ) + * h, + ( + D01[0] * K00[1] + + D01[1] * K01[1] + + D01[2] * K02[1] + + D01[3] * K03[1] + + D01[4] * K04[1] + + D01[5] * K05[1] + + D01[6] * K06[1] + + D01[7] * K07[1] + + D01[8] * K08[1] + + D01[9] * K09[1] + + D01[10] * K10[1] + + D01[11] * K11[1] + + D01[12] * K12[1] + + D01[13] * K13[1] + + D01[14] * K14[1] + + D01[15] * K15[1] + ) + * h, + ( + D01[0] * K00[2] + + D01[1] * K01[2] + + D01[2] * K02[2] + + D01[3] * K03[2] + + D01[4] * K04[2] + + D01[5] * K05[2] + + D01[6] * K06[2] + + D01[7] * K07[2] + + D01[8] * K08[2] + + D01[9] * K09[2] + + D01[10] * K10[2] + + D01[11] * K11[2] + + D01[12] * K12[2] + + D01[13] * K13[2] + + D01[14] * K14[2] + + D01[15] * K15[2] + ) + * h, + ( + D01[0] * K00[3] + + D01[1] * K01[3] + + D01[2] * K02[3] + + D01[3] * K03[3] + + D01[4] * K04[3] + + D01[5] * K05[3] + + D01[6] * K06[3] + + D01[7] * K07[3] + + D01[8] * K08[3] + + D01[9] * K09[3] + + D01[10] * K10[3] + + D01[11] * K11[3] + + D01[12] * K12[3] + + D01[13] * K13[3] + + D01[14] * K14[3] + + D01[15] * K15[3] + ) + * h, + ( + D01[0] * K00[4] + + D01[1] * K01[4] + + D01[2] * K02[4] + + D01[3] * K03[4] + + D01[4] * K04[4] + + D01[5] * K05[4] + + D01[6] * K06[4] + + D01[7] * K07[4] + + D01[8] * K08[4] + + D01[9] * K09[4] + + D01[10] * K10[4] + + D01[11] * K11[4] + + D01[12] * K12[4] + + D01[13] * K13[4] + + D01[14] * K14[4] + + D01[15] * K15[4] + ) + * h, + ( + D01[0] * K00[5] + + D01[1] * K01[5] + + D01[2] * K02[5] + + D01[3] * K03[5] + + D01[4] * K04[5] + + D01[5] * K05[5] + + D01[6] * K06[5] + + D01[7] * K07[5] + + D01[8] * K08[5] + + D01[9] * K09[5] + + D01[10] * K10[5] + + D01[11] * K11[5] + + D01[12] * K12[5] + + D01[13] * K13[5] + + D01[14] * K14[5] + + D01[15] * K15[5] + ) + * h, + ) + + F05 = ( + ( + D02[0] * K00[0] + + D02[1] * K01[0] + + D02[2] * K02[0] + + D02[3] * K03[0] + + D02[4] * K04[0] + + D02[5] * K05[0] + + D02[6] * K06[0] + + D02[7] * K07[0] + + D02[8] * K08[0] + + D02[9] * K09[0] + + D02[10] * K10[0] + + D02[11] * K11[0] + + D02[12] * K12[0] + + D02[13] * K13[0] + + D02[14] * K14[0] + + D02[15] * K15[0] + ) + * h, + ( + D02[0] * K00[1] + + D02[1] * K01[1] + + D02[2] * K02[1] + + D02[3] * K03[1] + + D02[4] * K04[1] + + D02[5] * K05[1] + + D02[6] * K06[1] + + D02[7] * K07[1] + + D02[8] * K08[1] + + D02[9] * K09[1] + + D02[10] * K10[1] + + D02[11] * K11[1] + + D02[12] * K12[1] + + D02[13] * K13[1] + + D02[14] * K14[1] + + D02[15] * K15[1] + ) + * h, + ( + D02[0] * K00[2] + + D02[1] * K01[2] + + D02[2] * K02[2] + + D02[3] * K03[2] + + D02[4] * K04[2] + + D02[5] * K05[2] + + D02[6] * K06[2] + + D02[7] * K07[2] + + D02[8] * K08[2] + + D02[9] * K09[2] + + D02[10] * K10[2] + + D02[11] * K11[2] + + D02[12] * K12[2] + + D02[13] * K13[2] + + D02[14] * K14[2] + + D02[15] * K15[2] + ) + * h, + ( + D02[0] * K00[3] + + D02[1] * K01[3] + + D02[2] * K02[3] + + D02[3] * K03[3] + + D02[4] * K04[3] + + D02[5] * K05[3] + + D02[6] * K06[3] + + D02[7] * K07[3] + + D02[8] * K08[3] + + D02[9] * K09[3] + + D02[10] * K10[3] + + D02[11] * K11[3] + + D02[12] * K12[3] + + D02[13] * K13[3] + + D02[14] * K14[3] + + D02[15] * K15[3] + ) + * h, + ( + D02[0] * K00[4] + + D02[1] * K01[4] + + D02[2] * K02[4] + + D02[3] * K03[4] + + D02[4] * K04[4] + + D02[5] * K05[4] + + D02[6] * K06[4] + + D02[7] * K07[4] + + D02[8] * K08[4] + + D02[9] * K09[4] + + D02[10] * K10[4] + + D02[11] * K11[4] + + D02[12] * K12[4] + + D02[13] * K13[4] + + D02[14] * K14[4] + + D02[15] * K15[4] + ) + * h, + ( + D02[0] * K00[5] + + D02[1] * K01[5] + + D02[2] * K02[5] + + D02[3] * K03[5] + + D02[4] * K04[5] + + D02[5] * K05[5] + + D02[6] * K06[5] + + D02[7] * K07[5] + + D02[8] * K08[5] + + D02[9] * K09[5] + + D02[10] * K10[5] + + D02[11] * K11[5] + + D02[12] * K12[5] + + D02[13] * K13[5] + + D02[14] * K14[5] + + D02[15] * K15[5] + ) + * h, + ) + + F06 = ( + ( + D03[0] * K00[0] + + D03[1] * K01[0] + + D03[2] * K02[0] + + D03[3] * K03[0] + + D03[4] * K04[0] + + D03[5] * K05[0] + + D03[6] * K06[0] + + D03[7] * K07[0] + + D03[8] * K08[0] + + D03[9] * K09[0] + + D03[10] * K10[0] + + D03[11] * K11[0] + + D03[12] * K12[0] + + D03[13] * K13[0] + + D03[14] * K14[0] + + D03[15] * K15[0] + ) + * h, + ( + D03[0] * K00[1] + + D03[1] * K01[1] + + D03[2] * K02[1] + + D03[3] * K03[1] + + D03[4] * K04[1] + + D03[5] * K05[1] + + D03[6] * K06[1] + + D03[7] * K07[1] + + D03[8] * K08[1] + + D03[9] * K09[1] + + D03[10] * K10[1] + + D03[11] * K11[1] + + D03[12] * K12[1] + + D03[13] * K13[1] + + D03[14] * K14[1] + + D03[15] * K15[1] + ) + * h, + ( + D03[0] * K00[2] + + D03[1] * K01[2] + + D03[2] * K02[2] + + D03[3] * K03[2] + + D03[4] * K04[2] + + D03[5] * K05[2] + + D03[6] * K06[2] + + D03[7] * K07[2] + + D03[8] * K08[2] + + D03[9] * K09[2] + + D03[10] * K10[2] + + D03[11] * K11[2] + + D03[12] * K12[2] + + D03[13] * K13[2] + + D03[14] * K14[2] + + D03[15] * K15[2] + ) + * h, + ( + D03[0] * K00[3] + + D03[1] * K01[3] + + D03[2] * K02[3] + + D03[3] * K03[3] + + D03[4] * K04[3] + + D03[5] * K05[3] + + D03[6] * K06[3] + + D03[7] * K07[3] + + D03[8] * K08[3] + + D03[9] * K09[3] + + D03[10] * K10[3] + + D03[11] * K11[3] + + D03[12] * K12[3] + + D03[13] * K13[3] + + D03[14] * K14[3] + + D03[15] * K15[3] + ) + * h, + ( + D03[0] * K00[4] + + D03[1] * K01[4] + + D03[2] * K02[4] + + D03[3] * K03[4] + + D03[4] * K04[4] + + D03[5] * K05[4] + + D03[6] * K06[4] + + D03[7] * K07[4] + + D03[8] * K08[4] + + D03[9] * K09[4] + + D03[10] * K10[4] + + D03[11] * K11[4] + + D03[12] * K12[4] + + D03[13] * K13[4] + + D03[14] * K14[4] + + D03[15] * K15[4] + ) + * h, + ( + D03[0] * K00[5] + + D03[1] * K01[5] + + D03[2] * K02[5] + + D03[3] * K03[5] + + D03[4] * K04[5] + + D03[5] * K05[5] + + D03[6] * K06[5] + + D03[7] * K07[5] + + D03[8] * K08[5] + + D03[9] * K09[5] + + D03[10] * K10[5] + + D03[11] * K11[5] + + D03[12] * K12[5] + + D03[13] * K13[5] + + D03[14] * K14[5] + + D03[15] * K15[5] + ) + * h, + ) return ( t_old, From f5645ae104a8180e707e8fb6e321e125fa8e7695 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 15:14:33 +0100 Subject: [PATCH 243/346] mv tuple decl --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 27 +++++++++++---------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index bb1662e02..ef289487a 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -45,6 +45,20 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): Ke = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) Ke[: N_STAGES + 1, :] = np.array(K_) + K00 = tuple(float_(number) for number in Ke[0, :]) + K01 = tuple(float_(number) for number in Ke[1, :]) + K02 = tuple(float_(number) for number in Ke[2, :]) + K03 = tuple(float_(number) for number in Ke[3, :]) + K04 = tuple(float_(number) for number in Ke[4, :]) + K05 = tuple(float_(number) for number in Ke[5, :]) + K06 = tuple(float_(number) for number in Ke[6, :]) + K07 = tuple(float_(number) for number in Ke[7, :]) + K08 = tuple(float_(number) for number in Ke[8, :]) + K09 = tuple(float_(number) for number in Ke[9, :]) + K10 = tuple(float_(number) for number in Ke[10, :]) + K11 = tuple(float_(number) for number in Ke[11, :]) + K12 = tuple(float_(number) for number in Ke[12, :]) + dy = np.dot(Ke[:13].T, np.array(A00)) * h rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) @@ -92,19 +106,6 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(fr, fr_old), h) ), *sub_VV_hf(mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(fv, fv_old), h)) - K00 = tuple(float_(number) for number in Ke[0, :]) - K01 = tuple(float_(number) for number in Ke[1, :]) - K02 = tuple(float_(number) for number in Ke[2, :]) - K03 = tuple(float_(number) for number in Ke[3, :]) - K04 = tuple(float_(number) for number in Ke[4, :]) - K05 = tuple(float_(number) for number in Ke[5, :]) - K06 = tuple(float_(number) for number in Ke[6, :]) - K07 = tuple(float_(number) for number in Ke[7, :]) - K08 = tuple(float_(number) for number in Ke[8, :]) - K09 = tuple(float_(number) for number in Ke[9, :]) - K10 = tuple(float_(number) for number in Ke[10, :]) - K11 = tuple(float_(number) for number in Ke[11, :]) - K12 = tuple(float_(number) for number in Ke[12, :]) K13 = tuple(float_(number) for number in Ke[13, :]) K14 = tuple(float_(number) for number in Ke[14, :]) K15 = tuple(float_(number) for number in Ke[15, :]) From 8ab797d5e86ea0b2577b71c1806d5d2f48ed5b50 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 15:16:43 +0100 Subject: [PATCH 244/346] passing k through --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index ef289487a..8536e0d3e 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -45,19 +45,7 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): Ke = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) Ke[: N_STAGES + 1, :] = np.array(K_) - K00 = tuple(float_(number) for number in Ke[0, :]) - K01 = tuple(float_(number) for number in Ke[1, :]) - K02 = tuple(float_(number) for number in Ke[2, :]) - K03 = tuple(float_(number) for number in Ke[3, :]) - K04 = tuple(float_(number) for number in Ke[4, :]) - K05 = tuple(float_(number) for number in Ke[5, :]) - K06 = tuple(float_(number) for number in Ke[6, :]) - K07 = tuple(float_(number) for number in Ke[7, :]) - K08 = tuple(float_(number) for number in Ke[8, :]) - K09 = tuple(float_(number) for number in Ke[9, :]) - K10 = tuple(float_(number) for number in Ke[10, :]) - K11 = tuple(float_(number) for number in Ke[11, :]) - K12 = tuple(float_(number) for number in Ke[12, :]) + K00, K01, K02, K03, K04, K05, K06, K07, K08, K09, K10, K11, K12 = K_ dy = np.dot(Ke[:13].T, np.array(A00)) * h rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) From 2836195a4f7dbbd4e3bbd0ded1b9bcbc43c60d7a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 15:22:06 +0100 Subject: [PATCH 245/346] rm np.dot --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 109 +++++++++++++++++++- 1 file changed, 104 insertions(+), 5 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index 8536e0d3e..c84f01e6e 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -47,16 +47,116 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): K00, K01, K02, K03, K04, K05, K06, K07, K08, K09, K10, K11, K12 = K_ - dy = np.dot(Ke[:13].T, np.array(A00)) * h - rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) - vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) + dr = ( + ( + K00[0] * A00[0] + + K01[0] * A00[1] + + K02[0] * A00[2] + + K03[0] * A00[3] + + K04[0] * A00[4] + + K05[0] * A00[5] + + K06[0] * A00[6] + + K07[0] * A00[7] + + K08[0] * A00[8] + + K09[0] * A00[9] + + K10[0] * A00[10] + + K11[0] * A00[11] + + K12[0] * A00[12] + ) + * h, + ( + K00[1] * A00[0] + + K01[1] * A00[1] + + K02[1] * A00[2] + + K03[1] * A00[3] + + K04[1] * A00[4] + + K05[1] * A00[5] + + K06[1] * A00[6] + + K07[1] * A00[7] + + K08[1] * A00[8] + + K09[1] * A00[9] + + K10[1] * A00[10] + + K11[1] * A00[11] + + K12[1] * A00[12] + ) + * h, + ( + K00[2] * A00[0] + + K01[2] * A00[1] + + K02[2] * A00[2] + + K03[2] * A00[3] + + K04[2] * A00[4] + + K05[2] * A00[5] + + K06[2] * A00[6] + + K07[2] * A00[7] + + K08[2] * A00[8] + + K09[2] * A00[9] + + K10[2] * A00[10] + + K11[2] * A00[11] + + K12[2] * A00[12] + ) + * h, + ) + dv = ( + ( + K00[3] * A00[0] + + K01[3] * A00[1] + + K02[3] * A00[2] + + K03[3] * A00[3] + + K04[3] * A00[4] + + K05[3] * A00[5] + + K06[3] * A00[6] + + K07[3] * A00[7] + + K08[3] * A00[8] + + K09[3] * A00[9] + + K10[3] * A00[10] + + K11[3] * A00[11] + + K12[3] * A00[12] + ) + * h, + ( + K00[4] * A00[0] + + K01[4] * A00[1] + + K02[4] * A00[2] + + K03[4] * A00[3] + + K04[4] * A00[4] + + K05[4] * A00[5] + + K06[4] * A00[6] + + K07[4] * A00[7] + + K08[4] * A00[8] + + K09[4] * A00[9] + + K10[4] * A00[10] + + K11[4] * A00[11] + + K12[4] * A00[12] + ) + * h, + ( + K00[5] * A00[0] + + K01[5] * A00[1] + + K02[5] * A00[2] + + K03[5] * A00[3] + + K04[5] * A00[4] + + K05[5] * A00[5] + + K06[5] * A00[6] + + K07[5] * A00[7] + + K08[5] * A00[8] + + K09[5] * A00[9] + + K10[5] * A00[10] + + K11[5] * A00[11] + + K12[5] * A00[12] + ) + * h, + ) + rr_ = add_VV_hf(rr_old, dr) + vv_ = add_VV_hf(vv_old, dv) rr_, vv_ = fun( t_old + C_EXTRA[0] * h, rr_, vv_, argk, ) - Ke[13] = np.array([*rr_, *vv_]) + K13 = *rr_, *vv_ + Ke[13, :] = np.array(K13) # TODO rm dy = np.dot(Ke[:14].T, np.array(A01)) * h rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) @@ -94,7 +194,6 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(fr, fr_old), h) ), *sub_VV_hf(mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(fv, fv_old), h)) - K13 = tuple(float_(number) for number in Ke[13, :]) K14 = tuple(float_(number) for number in Ke[14, :]) K15 = tuple(float_(number) for number in Ke[15, :]) From 39cec0a8ddfb3ad4ca0b54fd62371ff53f078182 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 15:25:13 +0100 Subject: [PATCH 246/346] rm np.dot --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 115 +++++++++++++++++++- 1 file changed, 110 insertions(+), 5 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index c84f01e6e..96d5c541b 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -158,16 +158,122 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): K13 = *rr_, *vv_ Ke[13, :] = np.array(K13) # TODO rm - dy = np.dot(Ke[:14].T, np.array(A01)) * h - rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) - vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) + dr = ( + ( + K00[0] * A01[0] + + K01[0] * A01[1] + + K02[0] * A01[2] + + K03[0] * A01[3] + + K04[0] * A01[4] + + K05[0] * A01[5] + + K06[0] * A01[6] + + K07[0] * A01[7] + + K08[0] * A01[8] + + K09[0] * A01[9] + + K10[0] * A01[10] + + K11[0] * A01[11] + + K12[0] * A01[12] + + K13[0] * A01[13] + ) + * h, + ( + K00[1] * A01[0] + + K01[1] * A01[1] + + K02[1] * A01[2] + + K03[1] * A01[3] + + K04[1] * A01[4] + + K05[1] * A01[5] + + K06[1] * A01[6] + + K07[1] * A01[7] + + K08[1] * A01[8] + + K09[1] * A01[9] + + K10[1] * A01[10] + + K11[1] * A01[11] + + K12[1] * A01[12] + + K13[1] * A01[13] + ) + * h, + ( + K00[2] * A01[0] + + K01[2] * A01[1] + + K02[2] * A01[2] + + K03[2] * A01[3] + + K04[2] * A01[4] + + K05[2] * A01[5] + + K06[2] * A01[6] + + K07[2] * A01[7] + + K08[2] * A01[8] + + K09[2] * A01[9] + + K10[2] * A01[10] + + K11[2] * A01[11] + + K12[2] * A01[12] + + K13[2] * A01[13] + ) + * h, + ) + dv = ( + ( + K00[3] * A01[0] + + K01[3] * A01[1] + + K02[3] * A01[2] + + K03[3] * A01[3] + + K04[3] * A01[4] + + K05[3] * A01[5] + + K06[3] * A01[6] + + K07[3] * A01[7] + + K08[3] * A01[8] + + K09[3] * A01[9] + + K10[3] * A01[10] + + K11[3] * A01[11] + + K12[3] * A01[12] + + K13[3] * A01[13] + ) + * h, + ( + K00[4] * A01[0] + + K01[4] * A01[1] + + K02[4] * A01[2] + + K03[4] * A01[3] + + K04[4] * A01[4] + + K05[4] * A01[5] + + K06[4] * A01[6] + + K07[4] * A01[7] + + K08[4] * A01[8] + + K09[4] * A01[9] + + K10[4] * A01[10] + + K11[4] * A01[11] + + K12[4] * A01[12] + + K13[4] * A01[13] + ) + * h, + ( + K00[5] * A01[0] + + K01[5] * A01[1] + + K02[5] * A01[2] + + K03[5] * A01[3] + + K04[5] * A01[4] + + K05[5] * A01[5] + + K06[5] * A01[6] + + K07[5] * A01[7] + + K08[5] * A01[8] + + K09[5] * A01[9] + + K10[5] * A01[10] + + K11[5] * A01[11] + + K12[5] * A01[12] + + K13[5] * A01[13] + ) + * h, + ) + rr_ = add_VV_hf(rr_old, dr) + vv_ = add_VV_hf(vv_old, dv) rr_, vv_ = fun( t_old + C_EXTRA[1] * h, rr_, vv_, argk, ) - Ke[14] = np.array([*rr_, *vv_]) + K14 = *rr_, *vv_ + Ke[14, :] = np.array(K14) # TODO rm dy = np.dot(Ke[:15].T, np.array(A02)) * h rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) @@ -194,7 +300,6 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(fr, fr_old), h) ), *sub_VV_hf(mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(fv, fv_old), h)) - K14 = tuple(float_(number) for number in Ke[14, :]) K15 = tuple(float_(number) for number in Ke[15, :]) F03 = ( From ce3a68ad436ce9583ed5e7b6e5552f8b86585773 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 15:27:59 +0100 Subject: [PATCH 247/346] rm np.dot --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 122 +++++++++++++++++++- 1 file changed, 116 insertions(+), 6 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index 96d5c541b..699e6432b 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -275,16 +275,128 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): K14 = *rr_, *vv_ Ke[14, :] = np.array(K14) # TODO rm - dy = np.dot(Ke[:15].T, np.array(A02)) * h - rr_ = add_VV_hf(rr_old, array_to_V_hf(dy[:3])) - vv_ = add_VV_hf(vv_old, array_to_V_hf(dy[3:])) + dr = ( + ( + K00[0] * A02[0] + + K01[0] * A02[1] + + K02[0] * A02[2] + + K03[0] * A02[3] + + K04[0] * A02[4] + + K05[0] * A02[5] + + K06[0] * A02[6] + + K07[0] * A02[7] + + K08[0] * A02[8] + + K09[0] * A02[9] + + K10[0] * A02[10] + + K11[0] * A02[11] + + K12[0] * A02[12] + + K13[0] * A02[13] + + K14[0] * A02[14] + ) + * h, + ( + K00[1] * A02[0] + + K01[1] * A02[1] + + K02[1] * A02[2] + + K03[1] * A02[3] + + K04[1] * A02[4] + + K05[1] * A02[5] + + K06[1] * A02[6] + + K07[1] * A02[7] + + K08[1] * A02[8] + + K09[1] * A02[9] + + K10[1] * A02[10] + + K11[1] * A02[11] + + K12[1] * A02[12] + + K13[1] * A02[13] + + K14[1] * A02[14] + ) + * h, + ( + K00[2] * A02[0] + + K01[2] * A02[1] + + K02[2] * A02[2] + + K03[2] * A02[3] + + K04[2] * A02[4] + + K05[2] * A02[5] + + K06[2] * A02[6] + + K07[2] * A02[7] + + K08[2] * A02[8] + + K09[2] * A02[9] + + K10[2] * A02[10] + + K11[2] * A02[11] + + K12[2] * A02[12] + + K13[2] * A02[13] + + K14[2] * A02[14] + ) + * h, + ) + dv = ( + ( + K00[3] * A02[0] + + K01[3] * A02[1] + + K02[3] * A02[2] + + K03[3] * A02[3] + + K04[3] * A02[4] + + K05[3] * A02[5] + + K06[3] * A02[6] + + K07[3] * A02[7] + + K08[3] * A02[8] + + K09[3] * A02[9] + + K10[3] * A02[10] + + K11[3] * A02[11] + + K12[3] * A02[12] + + K13[3] * A02[13] + + K14[3] * A02[14] + ) + * h, + ( + K00[4] * A02[0] + + K01[4] * A02[1] + + K02[4] * A02[2] + + K03[4] * A02[3] + + K04[4] * A02[4] + + K05[4] * A02[5] + + K06[4] * A02[6] + + K07[4] * A02[7] + + K08[4] * A02[8] + + K09[4] * A02[9] + + K10[4] * A02[10] + + K11[4] * A02[11] + + K12[4] * A02[12] + + K13[4] * A02[13] + + K14[4] * A02[14] + ) + * h, + ( + K00[5] * A02[0] + + K01[5] * A02[1] + + K02[5] * A02[2] + + K03[5] * A02[3] + + K04[5] * A02[4] + + K05[5] * A02[5] + + K06[5] * A02[6] + + K07[5] * A02[7] + + K08[5] * A02[8] + + K09[5] * A02[9] + + K10[5] * A02[10] + + K11[5] * A02[11] + + K12[5] * A02[12] + + K13[5] * A02[13] + + K14[5] * A02[14] + ) + * h, + ) + rr_ = add_VV_hf(rr_old, dr) + vv_ = add_VV_hf(vv_old, dv) rr_, vv_ = fun( t_old + C_EXTRA[2] * h, rr_, vv_, argk, ) - Ke[15] = np.array([*rr_, *vv_]) + K15 = *rr_, *vv_ + Ke[15, :] = np.array(K15) # TODO rm fr_old = array_to_V_hf(Ke[0, :3]) fv_old = array_to_V_hf(Ke[0, 3:]) @@ -300,8 +412,6 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): mul_Vs_hf(delta_rr, 2), mul_Vs_hf(add_VV_hf(fr, fr_old), h) ), *sub_VV_hf(mul_Vs_hf(delta_vv, 2), mul_Vs_hf(add_VV_hf(fv, fv_old), h)) - K15 = tuple(float_(number) for number in Ke[15, :]) - F03 = ( ( D00[0] * K00[0] From d5acb42871db6def62c4eabd113d4c3d293298cc Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 15:30:42 +0100 Subject: [PATCH 248/346] rm numpy dep --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 21 +++++---------------- 1 file changed, 5 insertions(+), 16 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index 699e6432b..16437427f 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -1,10 +1,4 @@ -import numpy as np - -from ._const import ( - N_RV, - N_STAGES, - N_STAGES_EXTENDED, -) +from ._const import N_STAGES from ._dop853_coefficients import A as _A, C as _C, D as _D from ..ieee754 import float_ from ..linalg import ( @@ -12,7 +6,8 @@ mul_Vs_hf, sub_VV_hf, ) -from ...jit import array_to_V_hf + +# from ...jit import hjit __all__ = [ "dense_output_hf", @@ -42,9 +37,6 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): assert t_old is not None assert t != t_old - Ke = np.empty((N_STAGES_EXTENDED, N_RV), dtype=float) - Ke[: N_STAGES + 1, :] = np.array(K_) - K00, K01, K02, K03, K04, K05, K06, K07, K08, K09, K10, K11, K12 = K_ dr = ( @@ -156,7 +148,6 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): argk, ) K13 = *rr_, *vv_ - Ke[13, :] = np.array(K13) # TODO rm dr = ( ( @@ -273,7 +264,6 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): argk, ) K14 = *rr_, *vv_ - Ke[14, :] = np.array(K14) # TODO rm dr = ( ( @@ -396,10 +386,9 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): argk, ) K15 = *rr_, *vv_ - Ke[15, :] = np.array(K15) # TODO rm - fr_old = array_to_V_hf(Ke[0, :3]) - fv_old = array_to_V_hf(Ke[0, 3:]) + fr_old = K00[:3] + fv_old = K00[3:] delta_rr = sub_VV_hf(rr, rr_old) delta_vv = sub_VV_hf(vv, vv_old) From f03c22bb5ccb19108853abe7055cafb690f28896 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 15:36:11 +0100 Subject: [PATCH 249/346] add compiler --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index 16437427f..9f11ea3f7 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -1,4 +1,4 @@ -from ._const import N_STAGES +from ._const import FSIG, KSIG, N_STAGES from ._dop853_coefficients import A as _A, C as _C, D as _D from ..ieee754 import float_ from ..linalg import ( @@ -6,8 +6,7 @@ mul_Vs_hf, sub_VV_hf, ) - -# from ...jit import hjit +from ...jit import hjit, DSIG __all__ = [ "dense_output_hf", @@ -24,7 +23,7 @@ D03 = tuple(float_(number) for number in _D[3, :]) -# TODO compile +@hjit(f"Tuple([f,f,V,V,{FSIG:s}])(F({DSIG:s}),f,f,f,f,V,V,V,V,V,V,{KSIG:s})") def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): """Compute a local interpolant over the last successful step. From 2c48d56743a3a94a0f9c8e89ca0cd50b06a693f7 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 15:40:19 +0100 Subject: [PATCH 250/346] name fix --- src/hapsira/core/math/ivp/_rkdenseoutput.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index 9f11ea3f7..984aadfa7 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -24,7 +24,7 @@ @hjit(f"Tuple([f,f,V,V,{FSIG:s}])(F({DSIG:s}),f,f,f,f,V,V,V,V,V,V,{KSIG:s})") -def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): +def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K): """Compute a local interpolant over the last successful step. Returns @@ -36,7 +36,7 @@ def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K_): assert t_old is not None assert t != t_old - K00, K01, K02, K03, K04, K05, K06, K07, K08, K09, K10, K11, K12 = K_ + K00, K01, K02, K03, K04, K05, K06, K07, K08, K09, K10, K11, K12 = K dr = ( ( From db265950a71dd99917b8221e4014771d43d2e9f1 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 17:00:45 +0100 Subject: [PATCH 251/346] mv module --- src/hapsira/core/math/ivp/{_rk.py => _rkcore.py} | 0 src/hapsira/core/math/ivp/_solve.py | 4 ++-- 2 files changed, 2 insertions(+), 2 deletions(-) rename src/hapsira/core/math/ivp/{_rk.py => _rkcore.py} (100%) diff --git a/src/hapsira/core/math/ivp/_rk.py b/src/hapsira/core/math/ivp/_rkcore.py similarity index 100% rename from src/hapsira/core/math/ivp/_rk.py rename to src/hapsira/core/math/ivp/_rkcore.py diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 01f3d2b62..3d9900043 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -3,9 +3,9 @@ import numpy as np from ._brentq import brentq_dense_hf, BRENTQ_CONVERGED, BRENTQ_MAXITER -from ._solution import OdeSolution -from ._rk import DOP853 +from ._rkcore import DOP853 from ._rkdenseoutput import dense_output_hf +from ._solution import OdeSolution from ..ieee754 import EPS From fe8fdc4093bda4f9ecbfc48114249f0bfaf337b2 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 21:19:40 +0100 Subject: [PATCH 252/346] dop853 class to functional via tuple --- src/hapsira/core/math/ivp/_rkcore.py | 370 +++++++++++++++++---------- src/hapsira/core/math/ivp/_solve.py | 60 +++-- 2 files changed, 278 insertions(+), 152 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkcore.py b/src/hapsira/core/math/ivp/_rkcore.py index d8afa0da2..b209326ef 100644 --- a/src/hapsira/core/math/ivp/_rkcore.py +++ b/src/hapsira/core/math/ivp/_rkcore.py @@ -1,150 +1,256 @@ +from math import nan from typing import Callable -import numpy as np - from ._const import ERROR_ESTIMATOR_ORDER from ._rkstepinit import select_initial_step_hf from ._rkstepimpl import step_impl_hf from ..ieee754 import EPS +from ..linalg import sign_hf __all__ = [ - "DOP853", + "dop853_init_hf", + "dop853_step_hf", + "DOP853_RUNNING", + "DOP853_FINISHED", + "DOP853_FAILED", + "DOP853_ARGK", + "DOP853_FR", + "DOP853_FUN", + "DOP853_FV", + "DOP853_H_PREVIOUS", + "DOP853_K", + "DOP853_RR", + "DOP853_RR_OLD", + "DOP853_STATUS", + "DOP853_T", + "DOP853_T_OLD", + "DOP853_VV", + "DOP853_VV_OLD", ] -class DOP853: +DOP853_RUNNING = 0 +DOP853_FINISHED = 1 +DOP853_FAILED = 2 + +DOP853_ARGK = 5 +DOP853_FR = 15 +DOP853_FUN = 4 +DOP853_FV = 16 +DOP853_H_PREVIOUS = 13 +DOP853_K = 9 +DOP853_RR = 1 +DOP853_RR_OLD = 10 +DOP853_STATUS = 14 +DOP853_T = 0 +DOP853_T_OLD = 12 +DOP853_VV = 2 +DOP853_VV_OLD = 11 + + +# TODO compile +def dop853_init_hf( + fun: Callable, + t0: float, + rr: tuple, + vv: tuple, + t_bound: float, + argk: float, + rtol: float, + atol: float, +): """ Explicit Runge-Kutta method of order 8. """ - def __init__( - self, - fun: Callable, - t0: float, - rr: tuple, - vv: tuple, - t_bound: float, - argk: float, - rtol: float, - atol: float, - ): - assert atol >= 0 - - if rtol < 100 * EPS: - rtol = 100 * EPS - - self.t = t0 - self.rr = rr - self.vv = vv - self.t_bound = t_bound - self.fun = fun - self.argk = argk - self.rtol = rtol - self.atol = atol - - self.direction = np.sign(t_bound - t0) if t_bound != t0 else 1 - - self.K = ( - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 0 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 1 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 2 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 3 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 4 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 5 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 6 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 7 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 8 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 9 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 10 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 11 - (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 12 - ) + assert atol >= 0 + + if rtol < 100 * EPS: + rtol = 100 * EPS + + direction = sign_hf(t_bound - t0) if t_bound != t0 else 1 + + K = ( + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 0 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 1 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 2 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 3 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 4 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 5 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 6 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 7 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 8 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 9 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 10 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 11 + (0.0, 0.0, 0.0, 0.0, 0.0, 0.0), # 12 + ) + + rr_old = (nan, nan, nan) + vv_old = (nan, nan, nan) + t_old = nan + h_previous = nan + + status = DOP853_RUNNING + + fr, fv = fun( + t0, + rr, + vv, + argk, + ) - self.rr_old = None - self.vv_old = None - self.t_old = None - self.h_previous = None - - self.status = "running" - - self.fr, self.fv = self.fun( - self.t, - self.rr, - self.vv, - self.argk, - ) # TODO call into hf - - self.h_abs = select_initial_step_hf( - self.fun, - self.t, - self.rr, - self.vv, - self.argk, - self.fr, - self.fv, - self.direction, - ERROR_ESTIMATOR_ORDER, - self.rtol, - self.atol, - ) # TODO call into hf - - def step(self): - """Perform one integration step. - - Returns - ------- - message : string or None - Report from the solver. Typically a reason for a failure if - `self.status` is 'failed' after the step was taken or None - otherwise. - """ - if self.status != "running": - raise RuntimeError("Attempt to step on a failed or finished " "solver.") - - if self.t == self.t_bound: - # Handle corner cases of empty solver or no integration. - self.t_old = self.t - self.t = self.t_bound - self.status = "finished" - return - - t = self.t - success, *rets = step_impl_hf( - self.fun, - self.argk, - self.t, - self.rr, - self.vv, - self.fr, - self.fv, - self.rtol, - self.atol, - self.direction, - self.h_abs, - self.t_bound, - self.K, + h_abs = select_initial_step_hf( + fun, + t0, + rr, + vv, + argk, + fr, + fv, + direction, + ERROR_ESTIMATOR_ORDER, + rtol, + atol, + ) + + return ( + t0, # 0 -> t + rr, # 1 + vv, # 2 + t_bound, # 3 + fun, # 4 + argk, # 5 + rtol, # 6 + atol, # 7 + direction, # 8 + K, # 9 + rr_old, # 10 + vv_old, # 11 + t_old, # 12 + h_previous, # 13 + status, # 14 + fr, # 15 + fv, # 16 + h_abs, # 17 + ) + + +# TODO compile +def dop853_step_hf( + t, + rr, + vv, + t_bound, + fun, + argk, + rtol, + atol, + direction, + K, + rr_old, + vv_old, + t_old, + h_previous, + status, + fr, + fv, + h_abs, +): + """Perform one integration step. + + Returns + ------- + message : string or None + Report from the solver. Typically a reason for a failure if + `self.status` is 'failed' after the step was taken or None + otherwise. + """ + + if status != DOP853_RUNNING: + raise RuntimeError("Attempt to step on a failed or finished " "solver.") + + if t == t_bound: + # Handle corner cases of empty solver or no integration. + t_old = t + t = t_bound + status = DOP853_FINISHED + return ( + t, + rr, + vv, + t_bound, + fun, + argk, + rtol, + atol, + direction, + K, + rr_old, + vv_old, + t_old, + h_previous, + status, + fr, + fv, + h_abs, ) - if success: - self.rr_old = self.rr - self.vv_old = self.vv - ( - self.h_previous, - self.t, - self.rr, - self.vv, - self.h_abs, - self.fr, - self.fv, - self.K, - ) = rets - - if not success: - self.status = "failed" - return - - self.t_old = t - if self.direction * (self.t - self.t_bound) < 0: - return - - self.status = "finished" + t_tmp = t + success, *rets = step_impl_hf( + fun, + argk, + t, + rr, + vv, + fr, + fv, + rtol, + atol, + direction, + h_abs, + t_bound, + K, + ) + + if success: + rr_old = rr + vv_old = vv + ( + h_previous, + t, + rr, + vv, + h_abs, + fr, + fv, + K, + ) = rets + + if not success: + status = DOP853_FAILED + else: + t_old = t_tmp + if not direction * (t - t_bound) < 0: + status = DOP853_FINISHED + + return ( + t, + rr, + vv, + t_bound, + fun, + argk, + rtol, + atol, + direction, + K, + rr_old, + vv_old, + t_old, + h_previous, + status, + fr, + fv, + h_abs, + ) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 3d9900043..206e74405 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -3,7 +3,25 @@ import numpy as np from ._brentq import brentq_dense_hf, BRENTQ_CONVERGED, BRENTQ_MAXITER -from ._rkcore import DOP853 +from ._rkcore import ( + dop853_init_hf, + dop853_step_hf, + DOP853_FINISHED, + DOP853_FAILED, + DOP853_ARGK, + DOP853_FR, + DOP853_FUN, + DOP853_FV, + DOP853_H_PREVIOUS, + DOP853_K, + DOP853_RR, + DOP853_RR_OLD, + DOP853_STATUS, + DOP853_T, + DOP853_T_OLD, + DOP853_VV, + DOP853_VV_OLD, +) from ._rkdenseoutput import dense_output_hf from ._solution import OdeSolution from ..ieee754 import EPS @@ -202,7 +220,7 @@ def solve_ivp( missed. """ - solver = DOP853(fun, t0, rr, vv, tf, argk, rtol, atol) + solver = dop853_init_hf(fun, t0, rr, vv, tf, argk, rtol, atol) ts = [t0] @@ -218,37 +236,39 @@ def solve_ivp( status = None while status is None: - solver.step() + solver = dop853_step_hf(*solver) - if solver.status == "finished": + if solver[DOP853_STATUS] == DOP853_FINISHED: status = 0 - elif solver.status == "failed": + elif solver[DOP853_STATUS] == DOP853_FAILED: status = -1 break - t_old = solver.t_old - t = solver.t + t_old = solver[DOP853_T_OLD] + t = solver[DOP853_T] sol = dense_output_hf( - solver.fun, - solver.argk, - solver.t_old, - solver.t, - solver.h_previous, - solver.rr, - solver.vv, - solver.rr_old, - solver.vv_old, - solver.fr, - solver.fv, - solver.K, + solver[DOP853_FUN], + solver[DOP853_ARGK], + solver[DOP853_T_OLD], + solver[DOP853_T], + solver[DOP853_H_PREVIOUS], + solver[DOP853_RR], + solver[DOP853_VV], + solver[DOP853_RR_OLD], + solver[DOP853_VV_OLD], + solver[DOP853_FR], + solver[DOP853_FV], + solver[DOP853_K], ) interpolants.append(sol) if events is not None: g_new = [] for event in events: - g_new.append(event.impl_hf(t, solver.rr, solver.vv, argk)) + g_new.append( + event.impl_hf(t, solver[DOP853_RR], solver[DOP853_VV], argk) + ) event.last_t_raw = t active_events = _find_active_events(g, g_new, event_dir) if active_events.size > 0: From f27565b4a932f0b353c6c167c523e5760367bdb8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 21:31:04 +0100 Subject: [PATCH 253/346] rk core compiled --- src/hapsira/core/math/ivp/_rkcore.py | 47 ++++++++++++++++++---------- 1 file changed, 31 insertions(+), 16 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkcore.py b/src/hapsira/core/math/ivp/_rkcore.py index b209326ef..9e1e67db1 100644 --- a/src/hapsira/core/math/ivp/_rkcore.py +++ b/src/hapsira/core/math/ivp/_rkcore.py @@ -1,11 +1,11 @@ from math import nan -from typing import Callable -from ._const import ERROR_ESTIMATOR_ORDER +from ._const import ERROR_ESTIMATOR_ORDER, KSIG from ._rkstepinit import select_initial_step_hf from ._rkstepimpl import step_impl_hf from ..ieee754 import EPS from ..linalg import sign_hf +from ...jit import hjit, DSIG __all__ = [ "dop853_init_hf", @@ -47,18 +47,32 @@ DOP853_VV = 2 DOP853_VV_OLD = 11 +_ = """ + t0, # 0 -> t + rr, # 1 + vv, # 2 + t_bound, # 3 + fun, # 4 + argk, # 5 + rtol, # 6 + atol, # 7 + direction, # 8 + K, # 9 + rr_old, # 10 + vv_old, # 11 + t_old, # 12 + h_previous, # 13 + status, # 14 + fr, # 15 + fv, # 16 + h_abs, # 17 +""" -# TODO compile -def dop853_init_hf( - fun: Callable, - t0: float, - rr: tuple, - vv: tuple, - t_bound: float, - argk: float, - rtol: float, - atol: float, -): +DOP853_SIG = f"f,V,V,f,F({DSIG}),f,f,f,f,{KSIG:s},V,V,f,f,f,V,V,f" + + +@hjit(f"Tuple([{DOP853_SIG:s}])(F({DSIG}),f,V,V,f,f,f,f)") +def dop853_init_hf(fun, t0, rr, vv, t_bound, argk, rtol, atol): """ Explicit Runge-Kutta method of order 8. """ @@ -136,7 +150,7 @@ def dop853_init_hf( ) -# TODO compile +@hjit(f"Tuple([{DOP853_SIG:s}])({DOP853_SIG:s})") def dop853_step_hf( t, rr, @@ -197,7 +211,7 @@ def dop853_step_hf( ) t_tmp = t - success, *rets = step_impl_hf( + rets = step_impl_hf( fun, argk, t, @@ -212,6 +226,7 @@ def dop853_step_hf( t_bound, K, ) + success = rets[0] if success: rr_old = rr @@ -225,7 +240,7 @@ def dop853_step_hf( fr, fv, K, - ) = rets + ) = rets[1:] if not success: status = DOP853_FAILED From 694a1e64a9c96e93871fde0b9ace739c1aa4a5b0 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 21:32:30 +0100 Subject: [PATCH 254/346] cleanup --- src/hapsira/core/math/ivp/_rkcore.py | 21 --------------------- 1 file changed, 21 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkcore.py b/src/hapsira/core/math/ivp/_rkcore.py index 9e1e67db1..12ceb83eb 100644 --- a/src/hapsira/core/math/ivp/_rkcore.py +++ b/src/hapsira/core/math/ivp/_rkcore.py @@ -47,27 +47,6 @@ DOP853_VV = 2 DOP853_VV_OLD = 11 -_ = """ - t0, # 0 -> t - rr, # 1 - vv, # 2 - t_bound, # 3 - fun, # 4 - argk, # 5 - rtol, # 6 - atol, # 7 - direction, # 8 - K, # 9 - rr_old, # 10 - vv_old, # 11 - t_old, # 12 - h_previous, # 13 - status, # 14 - fr, # 15 - fv, # 16 - h_abs, # 17 -""" - DOP853_SIG = f"f,V,V,f,F({DSIG}),f,f,f,f,{KSIG:s},V,V,f,f,f,V,V,f" From cb6a3a69101d41d83ca1cfbd3b4fb862e077c050 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 21:34:42 +0100 Subject: [PATCH 255/346] cleanup --- src/hapsira/core/math/ivp/_solve.py | 28 ++-------------------------- 1 file changed, 2 insertions(+), 26 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 206e74405..92b0b98e0 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -192,32 +192,8 @@ def solve_ivp( atol: float, events: Optional[List[Callable]] = None, ) -> Tuple[OdeSolution, bool]: - """Solve an initial value problem for a system of ODEs. - - Parameters - ---------- - fun : callable - Right-hand side of the system: the time derivative of the state ``y`` - at time ``t``. The calling signature is ``fun(t, y)``, where ``t`` is a - scalar and ``y`` is an ndarray with ``len(y) = len(y0)``. ``fun`` must - return an array of the same shape as ``y``. See `vectorized` for more - information. - t_span : 2-member sequence - Interval of integration (t0, tf). The solver starts with t=t0 and - integrates until it reaches t=tf. Both t0 and tf must be floats - or values interpretable by the float conversion function. - y0 : array_like, shape (n,) - Initial state. For problems in the complex domain, pass `y0` with a - complex data type (even if the initial value is purely real). - events : callable, or list of callables, optional - Events to track. If None (default), no events will be tracked. - Each event occurs at the zeros of a continuous function of time and - state. Each function must have the signature ``event(t, y)`` and return - a float. The solver will find an accurate value of `t` at which - ``event(t, y(t)) = 0`` using a root-finding algorithm. By default, all - zeros will be found. The solver looks for a sign change over each step, - so if multiple zero crossings occur within one step, events may be - missed. + """ + Solve an initial value problem for a system of ODEs. """ solver = dop853_init_hf(fun, t0, rr, vv, tf, argk, rtol, atol) From df53a1a3df926c2156b6f847bd15440c063e96f4 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 22:11:35 +0100 Subject: [PATCH 256/346] cleanup --- src/hapsira/core/math/ivp/_solve.py | 74 ++++++++--------------- src/hapsira/core/propagation/cowell.py | 11 ++-- src/hapsira/twobody/propagation/cowell.py | 2 +- 3 files changed, 32 insertions(+), 55 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 92b0b98e0..054ed2564 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -1,4 +1,4 @@ -from typing import Callable, List, Optional, Tuple +from typing import Callable, List, Tuple import numpy as np @@ -33,7 +33,11 @@ def _solve_event_equation( - event: Callable, sol: Callable, t_old: float, t: float, argk: float + event: Callable, + interpolant: Callable, + t_old: float, + t: float, + argk: float, ) -> float: """Solve an equation corresponding to an ODE event. @@ -65,7 +69,7 @@ def _solve_event_equation( 4 * EPS, 4 * EPS, BRENTQ_MAXITER, - *sol, + *interpolant, argk, ) event.last_t_raw = last_t @@ -74,10 +78,10 @@ def _solve_event_equation( def _handle_events( - sol, + interpolant, events: List[Callable], active_events, - is_terminal, + terminals, t_old: float, t: float, argk: float, @@ -93,7 +97,7 @@ def _handle_events( Event functions with signatures ``event(t, y)``. active_events : ndarray Indices of events which occurred. - is_terminal : ndarray, shape (n_events,) + terminals : ndarray, shape (n_events,) Which events are terminal. t_old, t : float Previous and new values of time. @@ -109,20 +113,20 @@ def _handle_events( Whether a terminal event occurred. """ roots = [ - _solve_event_equation(events[event_index], sol, t_old, t, argk) + _solve_event_equation(events[event_index], interpolant, t_old, t, argk) for event_index in active_events ] roots = np.asarray(roots) - if np.any(is_terminal[active_events]): + if np.any(terminals[active_events]): if t > t_old: order = np.argsort(roots) else: order = np.argsort(-roots) active_events = active_events[order] roots = roots[order] - t = np.nonzero(is_terminal[active_events])[0][0] + t = np.nonzero(terminals[active_events])[0][0] active_events = active_events[: t + 1] roots = roots[: t + 1] terminate = True @@ -132,39 +136,14 @@ def _handle_events( return active_events, roots, terminate -def _prepare_events(events): - """Standardize event functions and extract is_terminal and direction.""" - if callable(events): - events = (events,) - - if events is not None: - is_terminal = np.empty(len(events), dtype=bool) - direction = np.empty(len(events)) - for i, event in enumerate(events): - try: - is_terminal[i] = event.terminal - except AttributeError: - is_terminal[i] = False - - try: - direction[i] = event.direction - except AttributeError: - direction[i] = 0 - else: - is_terminal = None - direction = None - - return events, is_terminal, direction - - -def _find_active_events(g, g_new, direction): +def _find_active_events(g, g_new, directions): """Find which event occurred during an integration step. Parameters ---------- g, g_new : array_like, shape (n_events,) Values of event functions at a current and next points. - direction : ndarray, shape (n_events,) + directions : ndarray, shape (n_events,) Event "direction" according to the definition in `solve_ivp`. Returns @@ -176,7 +155,7 @@ def _find_active_events(g, g_new, direction): up = (g <= 0) & (g_new >= 0) down = (g >= 0) & (g_new <= 0) either = up | down - mask = up & (direction > 0) | down & (direction < 0) | either & (direction == 0) + mask = up & (directions > 0) | down & (directions < 0) | either & (directions == 0) return np.nonzero(mask)[0] @@ -190,21 +169,20 @@ def solve_ivp( argk: float, rtol: float, atol: float, - events: Optional[List[Callable]] = None, + events: Tuple[Callable], ) -> Tuple[OdeSolution, bool]: """ Solve an initial value problem for a system of ODEs. """ solver = dop853_init_hf(fun, t0, rr, vv, tf, argk, rtol, atol) - ts = [t0] - interpolants = [] - events, is_terminal, event_dir = _prepare_events(events) + terminals = np.array([event.terminal for event in events]) + directions = np.array([event.direction for event in events]) - if events is not None: + if len(events) > 0: g = [] for event in events: g.append(event.impl_hf(t0, rr, vv, argk)) @@ -223,7 +201,7 @@ def solve_ivp( t_old = solver[DOP853_T_OLD] t = solver[DOP853_T] - sol = dense_output_hf( + interpolant = dense_output_hf( solver[DOP853_FUN], solver[DOP853_ARGK], solver[DOP853_T_OLD], @@ -237,22 +215,22 @@ def solve_ivp( solver[DOP853_FV], solver[DOP853_K], ) - interpolants.append(sol) + interpolants.append(interpolant) - if events is not None: + if len(events) > 0: g_new = [] for event in events: g_new.append( event.impl_hf(t, solver[DOP853_RR], solver[DOP853_VV], argk) ) event.last_t_raw = t - active_events = _find_active_events(g, g_new, event_dir) + active_events = _find_active_events(g, g_new, directions) if active_events.size > 0: _, roots, terminate = _handle_events( - sol, + interpolant, events, active_events, - is_terminal, + terminals, t_old, t, argk, diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index b0f2c96a3..0f93ffea2 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -8,7 +8,7 @@ ] -def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=None, f=func_twobody_hf): +def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody_hf): """ Scalar cowell @@ -34,21 +34,20 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=None, f=func_twobody_hf argk=k, rtol=rtol, atol=atol, - events=events, + events=tuple(events), ) if not success: raise RuntimeError("Integration failed") - if events is not None: + if len(events) > 0: # Collect only the terminal events terminal_events = [event for event in events if event.terminal] # If there are no terminal events, then the last time of integration is the # greatest one from the original array of propagation times - if terminal_events: + if len(terminal_events) > 0: # Filter the event which triggered first - last_t = min(event._last_t for event in terminal_events) - # FIXME: Here last_t has units, but tofs don't + last_t = min(event.last_t_raw for event in terminal_events) tofs = [tof for tof in tofs if tof < last_t] tofs.append(last_t) diff --git a/src/hapsira/twobody/propagation/cowell.py b/src/hapsira/twobody/propagation/cowell.py index dd98d7855..e03304aec 100644 --- a/src/hapsira/twobody/propagation/cowell.py +++ b/src/hapsira/twobody/propagation/cowell.py @@ -27,7 +27,7 @@ class CowellPropagator: PropagatorKind.ELLIPTIC | PropagatorKind.PARABOLIC | PropagatorKind.HYPERBOLIC ) - def __init__(self, rtol=1e-11, events=None, f=func_twobody_hf): + def __init__(self, rtol=1e-11, events=tuple(), f=func_twobody_hf): self._rtol = rtol self._events = events self._f = f From 2e19322517001c839d4a41c0f1c7a0c5f67bc950 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Jan 2024 22:24:24 +0100 Subject: [PATCH 257/346] cleanup --- src/hapsira/core/math/ivp/_solve.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 054ed2564..3f7a69bb1 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -127,13 +127,12 @@ def _handle_events( active_events = active_events[order] roots = roots[order] t = np.nonzero(terminals[active_events])[0][0] - active_events = active_events[: t + 1] roots = roots[: t + 1] terminate = True else: terminate = False - return active_events, roots, terminate + return roots, terminate def _find_active_events(g, g_new, directions): @@ -226,7 +225,7 @@ def solve_ivp( event.last_t_raw = t active_events = _find_active_events(g, g_new, directions) if active_events.size > 0: - _, roots, terminate = _handle_events( + roots, terminate = _handle_events( interpolant, events, active_events, From dc6b2b2ed522ac5979b652481c04445152da9dc9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 30 Jan 2024 12:58:42 +0100 Subject: [PATCH 258/346] scalar is active --- src/hapsira/core/math/ivp/_solve.py | 40 +++++++++++++++++------------ 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 3f7a69bb1..d9e518520 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -25,6 +25,7 @@ from ._rkdenseoutput import dense_output_hf from ._solution import OdeSolution from ..ieee754 import EPS +from ...jit import hjit __all__ = [ @@ -112,6 +113,9 @@ def _handle_events( terminate : bool Whether a terminal event occurred. """ + + active_events = np.array(active_events) + roots = [ _solve_event_equation(events[event_index], interpolant, t_old, t, argk) for event_index in active_events @@ -135,7 +139,8 @@ def _handle_events( return roots, terminate -def _find_active_events(g, g_new, directions): +@hjit("b1(f,f,f)") +def _is_active_hf(g_old, g_new, direction): """Find which event occurred during an integration step. Parameters @@ -150,13 +155,11 @@ def _find_active_events(g, g_new, directions): active_events : ndarray Indices of events which occurred during the step. """ - g, g_new = np.asarray(g), np.asarray(g_new) - up = (g <= 0) & (g_new >= 0) - down = (g >= 0) & (g_new <= 0) + up = (g_old <= 0) & (g_new >= 0) + down = (g_old >= 0) & (g_new <= 0) either = up | down - mask = up & (directions > 0) | down & (directions < 0) | either & (directions == 0) - - return np.nonzero(mask)[0] + active = up & (direction > 0) | down & (direction < 0) | either & (direction == 0) + return active def solve_ivp( @@ -179,12 +182,11 @@ def solve_ivp( interpolants = [] terminals = np.array([event.terminal for event in events]) - directions = np.array([event.direction for event in events]) if len(events) > 0: - g = [] + gs_old = [] for event in events: - g.append(event.impl_hf(t0, rr, vv, argk)) + gs_old.append(event.impl_hf(t0, rr, vv, argk)) event.last_t_raw = t0 status = None @@ -217,18 +219,24 @@ def solve_ivp( interpolants.append(interpolant) if len(events) > 0: - g_new = [] + gs_new = [] for event in events: - g_new.append( + gs_new.append( event.impl_hf(t, solver[DOP853_RR], solver[DOP853_VV], argk) ) event.last_t_raw = t - active_events = _find_active_events(g, g_new, directions) - if active_events.size > 0: + + actives = [ + _is_active_hf(g_old, g_new, event.direction) + for g_old, g_new, event in zip(gs_old, gs_new, events) + ] + actives = [idx for idx, active in enumerate(actives) if active] + + if len(actives) > 0: roots, terminate = _handle_events( interpolant, events, - active_events, + actives, terminals, t_old, t, @@ -237,7 +245,7 @@ def solve_ivp( if terminate: status = 1 t = roots[-1] - g = g_new + gs_old = gs_new ts.append(t) From 7ffff095003b6235e9258f5f33ee2e0df0ced6d4 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 30 Jan 2024 12:59:56 +0100 Subject: [PATCH 259/346] fix name --- src/hapsira/core/math/ivp/_solve.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index d9e518520..fec9e289b 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -140,7 +140,7 @@ def _handle_events( @hjit("b1(f,f,f)") -def _is_active_hf(g_old, g_new, direction): +def _event_is_active_hf(g_old, g_new, direction): """Find which event occurred during an integration step. Parameters @@ -227,7 +227,7 @@ def solve_ivp( event.last_t_raw = t actives = [ - _is_active_hf(g_old, g_new, event.direction) + _event_is_active_hf(g_old, g_new, event.direction) for g_old, g_new, event in zip(gs_old, gs_new, events) ] actives = [idx for idx, active in enumerate(actives) if active] From 2145760cadfeb1bf4381e2696d97a373dd6e5f06 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 30 Jan 2024 14:07:14 +0100 Subject: [PATCH 260/346] allow to force inline --- src/hapsira/core/jit.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 0abeda00d..913d788b3 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -97,6 +97,11 @@ def hjit(*args, **kwargs) -> Callable: if len(args) > 0 and isinstance(args[0], str): args = _parse_signatures(args[0]), *args[1:] + try: + inline = kwargs.pop("inline") + except KeyError: + inline = settings["INLINE"].value + def wrapper(inner_func: Callable) -> Callable: """ Applies JIT @@ -106,14 +111,14 @@ def wrapper(inner_func: Callable) -> Callable: wjit = cuda.jit cfg = dict( device=True, - inline=settings["INLINE"].value, + inline=inline, cache=settings["CACHE"].value, ) else: wjit = nb.jit cfg = dict( nopython=settings["NOPYTHON"].value, - inline="always" if settings["INLINE"].value else "never", + inline="always" if inline else "never", cache=settings["CACHE"].value, ) cfg.update(kwargs) From 94f6dca0e5f39876f29698ad130bb10424794c8b Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 30 Jan 2024 14:07:41 +0100 Subject: [PATCH 261/346] simplify roots --- src/hapsira/core/math/ivp/_solve.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index fec9e289b..2d3a1690d 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -136,7 +136,7 @@ def _handle_events( else: terminate = False - return roots, terminate + return roots[-1], terminate @hjit("b1(f,f,f)") @@ -233,7 +233,7 @@ def solve_ivp( actives = [idx for idx, active in enumerate(actives) if active] if len(actives) > 0: - roots, terminate = _handle_events( + root, terminate = _handle_events( interpolant, events, actives, @@ -244,7 +244,7 @@ def solve_ivp( ) if terminate: status = 1 - t = roots[-1] + t = root gs_old = gs_new ts.append(t) From 52bdf1d5a4f841ed02e6132dcaf56f89bd076a36 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 30 Jan 2024 14:08:05 +0100 Subject: [PATCH 262/346] force inline --- src/hapsira/core/angles.py | 38 ++++++++++++++++----------------- src/hapsira/core/math/linalg.py | 36 +++++++++++++++---------------- 2 files changed, 37 insertions(+), 37 deletions(-) diff --git a/src/hapsira/core/angles.py b/src/hapsira/core/angles.py index 33bc04735..92b44c354 100644 --- a/src/hapsira/core/angles.py +++ b/src/hapsira/core/angles.py @@ -54,7 +54,7 @@ ] -@hjit("f(f,f)") +@hjit("f(f,f)", inline=True) def E_to_M_hf(E, ecc): r"""Mean anomaly from eccentric anomaly. @@ -99,7 +99,7 @@ def E_to_M_vf(E, ecc): return E_to_M_hf(E, ecc) -@hjit("f(f,f)") +@hjit("f(f,f)", inline=True) def F_to_M_hf(F, ecc): r"""Mean anomaly from hyperbolic anomaly. @@ -140,27 +140,27 @@ def F_to_M_vf(F, ecc): return F_to_M_hf(F, ecc) -@hjit("f(f,f,f)") +@hjit("f(f,f,f)", inline=True) def kepler_equation_hf(E, M, ecc): return E_to_M_hf(E, ecc) - M -@hjit("f(f,f,f)") +@hjit("f(f,f,f)", inline=True) def kepler_equation_prime_hf(E, M, ecc): return 1 - ecc * cos(E) -@hjit("f(f,f,f)") +@hjit("f(f,f,f)", inline=True) def kepler_equation_hyper_hf(F, M, ecc): return F_to_M_hf(F, ecc) - M -@hjit("f(f,f,f)") +@hjit("f(f,f,f)", inline=True) def kepler_equation_prime_hyper_hf(F, M, ecc): return ecc * cosh(F) - 1 -@hjit("f(f,f,f,f,i8)") +@hjit("f(f,f,f,f,i8)", inline=True) def _newton_elliptic_hf(p0, M, ecc, tol, maxiter): for _ in range(maxiter): fval = kepler_equation_hf(p0, M, ecc) @@ -173,7 +173,7 @@ def _newton_elliptic_hf(p0, M, ecc, tol, maxiter): return nan -@hjit("f(f,f,f,f,i8)") +@hjit("f(f,f,f,f,i8)", inline=True) def _newton_hyperbolic_hf(p0, M, ecc, tol, maxiter): for _ in range(maxiter): fval = kepler_equation_hyper_hf(p0, M, ecc) @@ -186,7 +186,7 @@ def _newton_hyperbolic_hf(p0, M, ecc, tol, maxiter): return nan -@hjit("f(f)") +@hjit("f(f)", inline=True) def D_to_nu_hf(D): r"""True anomaly from parabolic anomaly. @@ -221,7 +221,7 @@ def D_to_nu_vf(D): return D_to_nu_hf(D) -@hjit("f(f)") +@hjit("f(f)", inline=True) def nu_to_D_hf(nu): r"""Parabolic anomaly from true anomaly. @@ -282,7 +282,7 @@ def nu_to_D_vf(nu): return nu_to_D_hf(nu) -@hjit("f(f,f)") +@hjit("f(f,f)", inline=True) def nu_to_E_hf(nu, ecc): r"""Eccentric anomaly from true anomaly. @@ -327,7 +327,7 @@ def nu_to_E_vf(nu, ecc): return nu_to_E_hf(nu, ecc) -@hjit("f(f,f)") +@hjit("f(f,f)", inline=True) def nu_to_F_hf(nu, ecc): r"""Hyperbolic anomaly from true anomaly. @@ -372,7 +372,7 @@ def nu_to_F_vf(nu, ecc): return nu_to_F_hf(nu, ecc) -@hjit("f(f,f)") +@hjit("f(f,f)", inline=True) def E_to_nu_hf(E, ecc): r"""True anomaly from eccentric anomaly. @@ -417,7 +417,7 @@ def E_to_nu_vf(E, ecc): return E_to_nu_hf(E, ecc) -@hjit("f(f,f)") +@hjit("f(f,f)", inline=True) def F_to_nu_hf(F, ecc): r"""True anomaly from hyperbolic anomaly. @@ -455,7 +455,7 @@ def F_to_nu_vf(F, ecc): return F_to_nu_hf(F, ecc) -@hjit("f(f,f)") +@hjit("f(f,f)", inline=True) def M_to_E_hf(M, ecc): """Eccentric anomaly from mean anomaly. @@ -495,7 +495,7 @@ def M_to_E_vf(M, ecc): return M_to_E_hf(M, ecc) -@hjit("f(f,f)") +@hjit("f(f,f)", inline=True) def M_to_F_hf(M, ecc): """Hyperbolic anomaly from mean anomaly. @@ -530,7 +530,7 @@ def M_to_F_vf(M, ecc): return M_to_F_hf(M, ecc) -@hjit("f(f)") +@hjit("f(f)", inline=True) def M_to_D_hf(M): """Parabolic anomaly from mean anomaly. @@ -564,7 +564,7 @@ def M_to_D_vf(M): return M_to_D_hf(M) -@hjit("f(f)") +@hjit("f(f)", inline=True) def D_to_M_hf(D): r"""Mean anomaly from parabolic anomaly. @@ -606,7 +606,7 @@ def D_to_M_vf(D): return D_to_M_hf(D) -@hjit("f(f,f)") +@hjit("f(f,f)", inline=True) def fp_angle_hf(nu, ecc): r"""Returns the flight path angle. diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 226d069f0..7855d82c6 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -26,22 +26,22 @@ ] -@hjit("V(V)") +@hjit("V(V)", inline=True) def abs_V_hf(x): return abs(x[0]), abs(x[1]), abs(x[2]) -@hjit("V(V,f)") +@hjit("V(V,f)", inline=True) def add_Vs_hf(a, b): return a[0] + b, a[1] + b, a[2] + b -@hjit("V(V,V)") +@hjit("V(V,V)", inline=True) def add_VV_hf(a, b): return a[0] + b[0], a[1] + b[1], a[2] + b[2] -@hjit("V(V,V)") +@hjit("V(V,V)", inline=True) def cross_VV_hf(a, b): return ( a[1] * b[2] - a[2] * b[1], @@ -50,7 +50,7 @@ def cross_VV_hf(a, b): ) -@hjit("f(f,f)") +@hjit("f(f,f)", inline=True) def div_ss_hf(a, b): """ Similar to np.divide @@ -60,17 +60,17 @@ def div_ss_hf(a, b): return a / b -@hjit("V(V,V)") +@hjit("V(V,V)", inline=True) def div_VV_hf(x, y): return x[0] / y[0], x[1] / y[1], x[2] / y[2] -@hjit("V(V,f)") +@hjit("V(V,f)", inline=True) def div_Vs_hf(v, s): return v[0] / s, v[1] / s, v[2] / s -@hjit("M(M,M)") +@hjit("M(M,M)", inline=True) def matmul_MM_hf(a, b): return ( ( @@ -91,7 +91,7 @@ def matmul_MM_hf(a, b): ) -@hjit("V(V,M)") +@hjit("V(V,M)", inline=True) def matmul_VM_hf(a, b): return ( a[0] * b[0][0] + a[1] * b[1][0] + a[2] * b[2][0], @@ -100,12 +100,12 @@ def matmul_VM_hf(a, b): ) -@hjit("f(V,V)") +@hjit("f(V,V)", inline=True) def matmul_VV_hf(a, b): return a[0] * b[0] + a[1] * b[1] + a[2] * b[2] -@hjit("V(V,V)") +@hjit("V(V,V)", inline=True) def max_VV_hf(x, y): return ( x[0] if x[0] > y[0] else y[0], @@ -114,12 +114,12 @@ def max_VV_hf(x, y): ) -@hjit("V(V,f)") +@hjit("V(V,f)", inline=True) def mul_Vs_hf(v, s): return v[0] * s, v[1] * s, v[2] * s -@hjit("V(V,V)") +@hjit("V(V,V)", inline=True) def mul_VV_hf(a, b): return a[0] * b[0], a[1] * b[1], a[2] * b[2] @@ -131,7 +131,7 @@ def nextafter_hf(x, direction): return x - EPS -@hjit("f(V)") +@hjit("f(V)", inline=True) def norm_V_hf(a): return sqrt(matmul_VV_hf(a, a)) @@ -142,12 +142,12 @@ def norm_V_vf(a, b, c): return norm_V_hf((a, b, c)) -@hjit("f(V,V)") +@hjit("f(V,V)", inline=True) def norm_VV_hf(x, y): return sqrt(x[0] ** 2 + x[1] ** 2 + x[2] ** 2 + y[0] ** 2 + y[1] ** 2 + y[2] ** 2) -@hjit("f(f)") +@hjit("f(f)", inline=True) def sign_hf(x): if x < 0.0: return -1.0 @@ -156,12 +156,12 @@ def sign_hf(x): return 1.0 # if x > 0 -@hjit("V(V,V)") +@hjit("V(V,V)", inline=True) def sub_VV_hf(va, vb): return va[0] - vb[0], va[1] - vb[1], va[2] - vb[2] -@hjit("M(M)") +@hjit("M(M)", inline=True) def transpose_M_hf(a): return ( (a[0][0], a[1][0], a[2][0]), From c1df0ec8b285db6f2cc7e91e1c3aa3139288bfe3 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 31 Jan 2024 16:30:22 +0100 Subject: [PATCH 263/346] times increase, except if they're zero --- src/hapsira/core/math/ivp/_solve.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 2d3a1690d..e6349bed2 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -247,6 +247,11 @@ def solve_ivp( t = root gs_old = gs_new + try: + assert ts[-1] <= t + except AssertionError: + assert ts[-1] == 0 + ts.append(t) return OdeSolution(np.array(ts), interpolants), status >= 0 From 6944557f8a6bef75d290ddf8c47d37785e656f41 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 31 Jan 2024 16:32:53 +0100 Subject: [PATCH 264/346] refine assertion --- src/hapsira/core/math/ivp/_solve.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index e6349bed2..44221732a 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -216,7 +216,6 @@ def solve_ivp( solver[DOP853_FV], solver[DOP853_K], ) - interpolants.append(interpolant) if len(events) > 0: gs_new = [] @@ -248,10 +247,11 @@ def solve_ivp( gs_old = gs_new try: - assert ts[-1] <= t + assert ts[-1] < t except AssertionError: assert ts[-1] == 0 + interpolants.append(interpolant) ts.append(t) return OdeSolution(np.array(ts), interpolants), status >= 0 From 4b39933b1a3014fd14d0837abe6a631be43cb588 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 31 Jan 2024 16:37:23 +0100 Subject: [PATCH 265/346] cleanup --- src/hapsira/core/math/ivp/_solve.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 44221732a..bd67a2894 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -246,11 +246,7 @@ def solve_ivp( t = root gs_old = gs_new - try: - assert ts[-1] < t - except AssertionError: - assert ts[-1] == 0 - + assert ts[-1] <= t interpolants.append(interpolant) ts.append(t) From daacb796188c03fd79e30cb63f448cdfbcaa993a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 1 Feb 2024 11:07:03 +0100 Subject: [PATCH 266/346] indicators --- src/hapsira/core/math/ivp/_solve.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index bd67a2894..b84ec8939 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -123,13 +123,16 @@ def _handle_events( roots = np.asarray(roots) - if np.any(terminals[active_events]): + if np.any(terminals[active_events]): # is at least one terminal event active? if t > t_old: - order = np.argsort(roots) + order = np.argsort(roots) # sort index mask based on roots else: + raise ValueError("not t > t_old", t, t_old) order = np.argsort(-roots) - active_events = active_events[order] - roots = roots[order] + + active_events = active_events[order] # sort active_events + roots = roots[order] # sort roots + t = np.nonzero(terminals[active_events])[0][0] roots = roots[: t + 1] terminate = True @@ -246,7 +249,8 @@ def solve_ivp( t = root gs_old = gs_new - assert ts[-1] <= t + if not ts[-1] <= t: + raise ValueError("not ts[-1] <= t", ts[-1], t) interpolants.append(interpolant) ts.append(t) From 1e45f9cc2f6a6e92319b0f6c00178449f190bb17 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 1 Feb 2024 11:28:35 +0100 Subject: [PATCH 267/346] rm ode solution class --- src/hapsira/core/math/ivp/_solution.py | 76 -------------------------- src/hapsira/core/math/ivp/_solve.py | 24 +++++++- 2 files changed, 21 insertions(+), 79 deletions(-) delete mode 100644 src/hapsira/core/math/ivp/_solution.py diff --git a/src/hapsira/core/math/ivp/_solution.py b/src/hapsira/core/math/ivp/_solution.py deleted file mode 100644 index 0ed38344a..000000000 --- a/src/hapsira/core/math/ivp/_solution.py +++ /dev/null @@ -1,76 +0,0 @@ -import numpy as np - -from ._rkdenseinterp import dense_interp_hf - - -class OdeSolution: - """Continuous ODE solution. - - It is organized as a collection of `DenseOutput` objects which represent - local interpolants. It provides an algorithm to select a right interpolant - for each given point. - - The interpolants cover the range between `t_min` and `t_max` (see - Attributes below). Evaluation outside this interval is not forbidden, but - the accuracy is not guaranteed. - - When evaluating at a breakpoint (one of the values in `ts`) a segment with - the lower index is selected. - - Parameters - ---------- - ts : array_like, shape (n_segments + 1,) - Time instants between which local interpolants are defined. Must - be strictly increasing or decreasing (zero segment with two points is - also allowed). - interpolants : list of DenseOutput with n_segments elements - Local interpolants. An i-th interpolant is assumed to be defined - between ``ts[i]`` and ``ts[i + 1]``. - - Attributes - ---------- - t_min, t_max : float - Time range of the interpolation. - """ - - def __init__(self, ts, interpolants): - ts = np.asarray(ts) - d = np.diff(ts) - # The first case covers integration on zero segment. - if not ((ts.size == 2 and ts[0] == ts[-1]) or np.all(d > 0) or np.all(d < 0)): - raise ValueError("`ts` must be strictly increasing or decreasing.") - - self.n_segments = len(interpolants) - if ts.shape != (self.n_segments + 1,): - raise ValueError("Numbers of time stamps and interpolants " "don't match.") - - self.ts = ts - self.interpolants = interpolants - - assert ts[-1] >= ts[0] - - self.t_min = ts[0] - self.t_max = ts[-1] - self.ts_sorted = ts - - def __call__(self, t): - """Evaluate the solution. - - Parameters - ---------- - t : float or array_like with shape (n_points,) - Points to evaluate at. - - Returns - ------- - y : ndarray, shape (n_states,) or (n_states, n_points) - Computed values. Shape depends on whether `t` is a scalar or a - 1-D array. - """ - # t = np.asarray(t) - - # assert t.ndim == 0 - - ind = np.searchsorted(self.ts_sorted, t, side="left") - segment = min(max(ind - 1, 0), self.n_segments - 1) - return dense_interp_hf(t, *self.interpolants[segment]) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index b84ec8939..4e6669819 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -22,8 +22,8 @@ DOP853_VV, DOP853_VV_OLD, ) +from ._rkdenseinterp import dense_interp_hf from ._rkdenseoutput import dense_output_hf -from ._solution import OdeSolution from ..ieee754 import EPS from ...jit import hjit @@ -175,7 +175,7 @@ def solve_ivp( rtol: float, atol: float, events: Tuple[Callable], -) -> Tuple[OdeSolution, bool]: +) -> Tuple[Callable, bool]: """ Solve an initial value problem for a system of ODEs. """ @@ -254,4 +254,22 @@ def solve_ivp( interpolants.append(interpolant) ts.append(t) - return OdeSolution(np.array(ts), interpolants), status >= 0 + assert len(ts) >= 2 + assert len(ts) == len(interpolants) + 1 + assert ( + (len(ts) == 2 and ts[0] == ts[1]) + or all(a - b > 0 for a, b in zip(ts[:-1], ts[1:])) + or all(b - a > 0 for a, b in zip(ts[:-1], ts[1:])) + ) + + def ode_solution( + t: float, + ) -> Tuple[Tuple[float, float, float], Tuple[float, float, float]]: + """ + Evaluate the solution + """ + idx = np.searchsorted(ts, t, side="left") + segment = min(max(idx - 1, 0), len(interpolants) - 1) + return dense_interp_hf(t, *interpolants[segment]) + + return ode_solution, status >= 0 From 9b07188575863f7e046c676b8e534f9429e5615d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 1 Feb 2024 14:57:17 +0100 Subject: [PATCH 268/346] assumptions --- src/hapsira/core/propagation/cowell.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 0f93ffea2..a9b14c826 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -24,6 +24,8 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody assert hasattr(f, "djit") # DEBUG check for compiler flag assert isinstance(rtol, float) + assert all(tof >= 0 for tof in tofs) + assert sorted(tofs) == list(tofs) sol, success = solve_ivp( f, From 9353d2fd394adc31ac39f8314bbe736a0d1f44a6 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 1 Feb 2024 14:57:59 +0100 Subject: [PATCH 269/346] add note --- src/hapsira/core/propagation/cowell.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index a9b14c826..13081bccf 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -20,6 +20,8 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody atol : float events : Optional[List[Event]] f : Callable + + Can be reversed: https://github.com/poliastro/poliastro/issues/1630 """ assert hasattr(f, "djit") # DEBUG check for compiler flag From 89296229d6b105bb9873e2c8a7307d72f93fe942 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 1 Feb 2024 17:38:22 +0100 Subject: [PATCH 270/346] fix link --- src/hapsira/twobody/propagation/cowell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/twobody/propagation/cowell.py b/src/hapsira/twobody/propagation/cowell.py index e03304aec..e7410e125 100644 --- a/src/hapsira/twobody/propagation/cowell.py +++ b/src/hapsira/twobody/propagation/cowell.py @@ -63,7 +63,7 @@ def propagate_many(self, state, tofs): ) # TODO: This should probably return a RVStateArray instead, - # see discussion at https://github.com/hapsira/hapsira/pull/1492 + # see discussion at https://github.com/poliastro/poliastro/pull/1492 return ( rrs << u.km, vvs << (u.km / u.s), From 81957b9847fe1b9d6eee81f70e3fcc9d7703e97a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 2 Feb 2024 12:52:58 +0100 Subject: [PATCH 271/346] name fix --- src/hapsira/core/math/ivp/__init__.py | 6 +++--- src/hapsira/core/math/ivp/_rkdenseinterp.py | 10 +++++----- src/hapsira/core/math/ivp/_rkdenseoutput.py | 4 ++-- src/hapsira/core/math/ivp/_solve.py | 8 ++++---- src/hapsira/twobody/events.py | 4 ++-- 5 files changed, 16 insertions(+), 16 deletions(-) diff --git a/src/hapsira/core/math/ivp/__init__.py b/src/hapsira/core/math/ivp/__init__.py index f918f2b4d..173e3926d 100644 --- a/src/hapsira/core/math/ivp/__init__.py +++ b/src/hapsira/core/math/ivp/__init__.py @@ -8,7 +8,7 @@ BRENTQ_MAXITER, brentq_hf, ) -from ._rkdenseinterp import dense_interp_brentq_hb, dense_interp_hf +from ._rkdenseinterp import dop853_dense_interp_brentq_hb, dop853_dense_interp_hf from ._solve import solve_ivp __all__ = [ @@ -20,7 +20,7 @@ "BRENTQ_RTOL", "BRENTQ_MAXITER", "brentq_hf", - "dense_interp_brentq_hb", - "dense_interp_hf", + "dop853_dense_interp_brentq_hb", + "dop853_dense_interp_hf", "solve_ivp", ] diff --git a/src/hapsira/core/math/ivp/_rkdenseinterp.py b/src/hapsira/core/math/ivp/_rkdenseinterp.py index 9c6aa442b..2d1e7288a 100644 --- a/src/hapsira/core/math/ivp/_rkdenseinterp.py +++ b/src/hapsira/core/math/ivp/_rkdenseinterp.py @@ -4,8 +4,8 @@ __all__ = [ - "dense_interp_brentq_hb", - "dense_interp_hf", + "dop853_dense_interp_brentq_hb", + "dop853_dense_interp_hf", "DENSE_SIG", ] @@ -14,7 +14,7 @@ @hjit(f"Tuple([V,V])(f,{DENSE_SIG:s})") -def dense_interp_hf(t, t_old, h, rr_old, vv_old, F): +def dop853_dense_interp_hf(t, t_old, h, rr_old, vv_old, F): """ Local interpolant over step made by an ODE solver. Evaluate the interpolant. @@ -74,10 +74,10 @@ def dense_interp_hf(t, t_old, h, rr_old, vv_old, F): return rr_new, vv_new -def dense_interp_brentq_hb(func): +def dop853_dense_interp_brentq_hb(func): @hjit(f"f(f,{DENSE_SIG:s},f)", cache=False) def event_wrapper(t, t_old, h, rr_old, vv_old, F, argk): - rr, vv = dense_interp_hf(t, t_old, h, rr_old, vv_old, F) + rr, vv = dop853_dense_interp_hf(t, t_old, h, rr_old, vv_old, F) return func(t, rr, vv, argk) return event_wrapper diff --git a/src/hapsira/core/math/ivp/_rkdenseoutput.py b/src/hapsira/core/math/ivp/_rkdenseoutput.py index 984aadfa7..c5ee39ef7 100644 --- a/src/hapsira/core/math/ivp/_rkdenseoutput.py +++ b/src/hapsira/core/math/ivp/_rkdenseoutput.py @@ -9,7 +9,7 @@ from ...jit import hjit, DSIG __all__ = [ - "dense_output_hf", + "dop853_dense_output_hf", ] @@ -24,7 +24,7 @@ @hjit(f"Tuple([f,f,V,V,{FSIG:s}])(F({DSIG:s}),f,f,f,f,V,V,V,V,V,V,{KSIG:s})") -def dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K): +def dop853_dense_output_hf(fun, argk, t_old, t, h, rr, vv, rr_old, vv_old, fr, fv, K): """Compute a local interpolant over the last successful step. Returns diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 4e6669819..247190ff9 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -22,8 +22,8 @@ DOP853_VV, DOP853_VV_OLD, ) -from ._rkdenseinterp import dense_interp_hf -from ._rkdenseoutput import dense_output_hf +from ._rkdenseinterp import dop853_dense_interp_hf +from ._rkdenseoutput import dop853_dense_output_hf from ..ieee754 import EPS from ...jit import hjit @@ -205,7 +205,7 @@ def solve_ivp( t_old = solver[DOP853_T_OLD] t = solver[DOP853_T] - interpolant = dense_output_hf( + interpolant = dop853_dense_output_hf( solver[DOP853_FUN], solver[DOP853_ARGK], solver[DOP853_T_OLD], @@ -270,6 +270,6 @@ def ode_solution( """ idx = np.searchsorted(ts, t, side="left") segment = min(max(idx - 1, 0), len(interpolants) - 1) - return dense_interp_hf(t, *interpolants[segment]) + return dop853_dense_interp_hf(t, *interpolants[segment]) return ode_solution, status >= 0 diff --git a/src/hapsira/twobody/events.py b/src/hapsira/twobody/events.py index 6b0783e3b..6b60ccdb0 100644 --- a/src/hapsira/twobody/events.py +++ b/src/hapsira/twobody/events.py @@ -6,7 +6,7 @@ from astropy.coordinates import get_body_barycentric_posvel from hapsira.core.jit import hjit -from hapsira.core.math.ivp import dense_interp_brentq_hb +from hapsira.core.math.ivp import dop853_dense_interp_brentq_hb from hapsira.core.math.linalg import mul_Vs_hf, norm_V_hf from hapsira.core.events import ( eclipse_function_hf, @@ -78,7 +78,7 @@ def impl_dense_hf(self) -> Callable: return self._impl_dense_hf def _wrap(self): - self._impl_dense_hf = dense_interp_brentq_hb(self._impl_hf) + self._impl_dense_hf = dop853_dense_interp_brentq_hb(self._impl_hf) class AltitudeCrossEvent(BaseEvent): From 5b1ea67b781e23388958dbc8e08332894a7b1c56 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 2 Feb 2024 19:30:41 +0100 Subject: [PATCH 272/346] notes --- src/hapsira/core/math/ivp/_solve.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 247190ff9..2ce02500f 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -184,6 +184,20 @@ def solve_ivp( ts = [t0] interpolants = [] + _ = """ + Event: + - impl_hf (callable) -> compiled tuple + - impl_dense_hf (callable) -> compiled tuple + - terminal (const) -> input array + - direction (const) -> input array + - is_active + - g_old + - g_new + - last_t -> output array + N events -> compiled const int? + for-loop??? + """ + terminals = np.array([event.terminal for event in events]) if len(events) > 0: From 4489db76559f32ccf1db8f6ed4c69e06c7c37f73 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 3 Feb 2024 17:27:08 +0100 Subject: [PATCH 273/346] rm numpy from handling events --- src/hapsira/core/math/ivp/_solve.py | 145 +++++++++++++++---------- src/hapsira/core/propagation/cowell.py | 41 ++++++- 2 files changed, 126 insertions(+), 60 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 2ce02500f..659635bad 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -1,3 +1,4 @@ +from math import nan, isnan from typing import Callable, List, Tuple import numpy as np @@ -80,9 +81,10 @@ def _solve_event_equation( def _handle_events( interpolant, - events: List[Callable], - active_events, - terminals, + event_impl_dense_hfs: List[Callable], + event_last_ts: np.ndarray, + event_actives: np.ndarray, + event_terminals: np.ndarray, t_old: float, t: float, argk: float, @@ -114,32 +116,49 @@ def _handle_events( Whether a terminal event occurred. """ - active_events = np.array(active_events) + assert np.any(event_actives) # nothing active - roots = [ - _solve_event_equation(events[event_index], interpolant, t_old, t, argk) - for event_index in active_events - ] + EVENTS = len(event_impl_dense_hfs) # TODO compile as const - roots = np.asarray(roots) + pivot = nan # set initial value + terminate = False - if np.any(terminals[active_events]): # is at least one terminal event active? - if t > t_old: - order = np.argsort(roots) # sort index mask based on roots - else: - raise ValueError("not t > t_old", t, t_old) - order = np.argsort(-roots) + for idx in range(EVENTS): + if not event_actives[idx]: + continue - active_events = active_events[order] # sort active_events - roots = roots[order] # sort roots + event_last_ts[idx], root, status = brentq_dense_hf( + event_impl_dense_hfs[idx], + t_old, + t, + 4 * EPS, + 4 * EPS, + BRENTQ_MAXITER, + *interpolant, + argk, + ) + assert status == BRENTQ_CONVERGED + + if event_terminals[idx]: + terminate = True + + if isnan(pivot): + pivot = root + continue + + if t > t_old: # smallest root of all active events + if root < pivot: + pivot = root + continue - t = np.nonzero(terminals[active_events])[0][0] - roots = roots[: t + 1] - terminate = True - else: - terminate = False + # largest root of all active events + if root > pivot: + pivot = root + raise ValueError("not t > t_old", t, t_old) # TODO remove - return roots[-1], terminate + assert not isnan(pivot) + + return pivot if terminate else nan, terminate @hjit("b1(f,f,f)") @@ -174,12 +193,21 @@ def solve_ivp( argk: float, rtol: float, atol: float, - events: Tuple[Callable], + event_impl_hfs: Tuple[Callable, ...], + event_impl_dense_hfs: Tuple[Callable, ...], + event_terminals: np.ndarray, + event_directions: np.ndarray, + event_actives: np.ndarray, + event_g_olds: np.ndarray, + event_g_news: np.ndarray, + event_last_ts: np.ndarray, ) -> Tuple[Callable, bool]: """ Solve an initial value problem for a system of ODEs. """ + EVENTS = len(event_impl_hfs) # TODO compile as const + solver = dop853_init_hf(fun, t0, rr, vv, tf, argk, rtol, atol) ts = [t0] interpolants = [] @@ -198,13 +226,9 @@ def solve_ivp( for-loop??? """ - terminals = np.array([event.terminal for event in events]) - - if len(events) > 0: - gs_old = [] - for event in events: - gs_old.append(event.impl_hf(t0, rr, vv, argk)) - event.last_t_raw = t0 + for idx in range(EVENTS): + event_g_olds[idx] = event_impl_hfs[idx](t0, rr, vv, argk) + event_last_ts[idx] = t0 status = None while status is None: @@ -234,34 +258,37 @@ def solve_ivp( solver[DOP853_K], ) - if len(events) > 0: - gs_new = [] - for event in events: - gs_new.append( - event.impl_hf(t, solver[DOP853_RR], solver[DOP853_VV], argk) - ) - event.last_t_raw = t - - actives = [ - _event_is_active_hf(g_old, g_new, event.direction) - for g_old, g_new, event in zip(gs_old, gs_new, events) - ] - actives = [idx for idx, active in enumerate(actives) if active] - - if len(actives) > 0: - root, terminate = _handle_events( - interpolant, - events, - actives, - terminals, - t_old, - t, - argk, - ) - if terminate: - status = 1 - t = root - gs_old = gs_new + at_least_one_active = False + for idx in range(EVENTS): + event_g_news[idx] = event_impl_hfs[idx]( + t, solver[DOP853_RR], solver[DOP853_VV], argk + ) + event_last_ts[idx] = t + event_actives[idx] = _event_is_active_hf( + event_g_olds[idx], + event_g_news[idx], + event_directions[idx], + ) + if event_actives[idx]: + at_least_one_active = True + + if at_least_one_active: + root, terminate = _handle_events( + interpolant, + event_impl_dense_hfs, # TODO + event_last_ts, + event_actives, # TODO + event_terminals, # TODO + t_old, + t, + argk, + ) + if terminate: + status = 1 + t = root + + for idx in range(EVENTS): + event_g_olds[idx] = event_g_news[idx] if not ts[-1] <= t: raise ValueError("not ts[-1] <= t", ts[-1], t) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 13081bccf..d62a73694 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -1,3 +1,5 @@ +import numpy as np + from ..jit import array_to_V_hf from ..math.ivp import solve_ivp from ..propagation.base import func_twobody_hf @@ -29,6 +31,33 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody assert all(tof >= 0 for tof in tofs) assert sorted(tofs) == list(tofs) + EVENTS = len(events) # TODO compile as const + + event_impl_hfs = tuple( + event.impl_hf for event in events + ) # TODO compile into kernel + event_impl_dense_hfs = tuple( + event.impl_dense_hf for event in events + ) # TODO compile into kernel + event_terminals = np.array( + [event.terminal for event in events], dtype=bool + ) # gufunc param static + event_directions = np.array( + [event.direction for event in events], dtype=float + ) # gufunc param static + event_actives = np.full( + (EVENTS,), fill_value=np.nan, dtype=bool + ) # gufunc param TODO reset to nan + event_g_olds = np.full( + (EVENTS,), fill_value=np.nan, dtype=float + ) # gufunc param TODO reset to nan + event_g_news = np.full( + (EVENTS,), fill_value=np.nan, dtype=float + ) # gufunc param TODO reset to nan + event_last_ts = np.full( + (EVENTS,), fill_value=np.nan, dtype=float + ) # gufunc param TODO reset to nan + sol, success = solve_ivp( f, 0.0, @@ -38,11 +67,21 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody argk=k, rtol=rtol, atol=atol, - events=tuple(events), + event_impl_hfs=event_impl_hfs, + event_impl_dense_hfs=event_impl_dense_hfs, + event_terminals=event_terminals, + event_directions=event_directions, + event_actives=event_actives, + event_g_olds=event_g_olds, + event_g_news=event_g_news, + event_last_ts=event_last_ts, ) if not success: raise RuntimeError("Integration failed") + for idx in range(EVENTS): + events[idx].last_t_raw = event_last_ts[idx] + if len(events) > 0: # Collect only the terminal events terminal_events = [event for event in events if event.terminal] From 4294b6b8b9c7611eab385003abd276477434e5cf Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 3 Feb 2024 17:35:45 +0100 Subject: [PATCH 274/346] cleanup --- src/hapsira/core/math/ivp/_solve.py | 45 ----------------------------- 1 file changed, 45 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 659635bad..2e226893c 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -34,51 +34,6 @@ ] -def _solve_event_equation( - event: Callable, - interpolant: Callable, - t_old: float, - t: float, - argk: float, -) -> float: - """Solve an equation corresponding to an ODE event. - - The equation is ``event(t, y(t)) = 0``, here ``y(t)`` is known from an - ODE solver using some sort of interpolation. It is solved by - `scipy.optimize.brentq` with xtol=atol=4*EPS. - - Parameters - ---------- - event : callable - Function ``event(t, y)``. - sol : callable - Function ``sol(t)`` which evaluates an ODE solution between `t_old` - and `t`. - t_old, t : float - Previous and new values of time. They will be used as a bracketing - interval. - - Returns - ------- - root : float - Found solution. - """ - - last_t, value, status = brentq_dense_hf( - event.impl_dense_hf, - t_old, - t, - 4 * EPS, - 4 * EPS, - BRENTQ_MAXITER, - *interpolant, - argk, - ) - event.last_t_raw = last_t - assert BRENTQ_CONVERGED == status - return value - - def _handle_events( interpolant, event_impl_dense_hfs: List[Callable], From c6916b5be633da645f2c3aaf8fb2560c285512ff Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 4 Feb 2024 12:49:23 +0100 Subject: [PATCH 275/346] workaround: https://github.com/numba/numba/issues/9420 --- src/hapsira/core/math/ivp/_brentq.py | 11 ++-- src/hapsira/core/math/ivp/_const.py | 2 + src/hapsira/core/math/ivp/_rkdenseinterp.py | 6 +- src/hapsira/core/math/ivp/_solve.py | 73 +++++++++++++++++++-- 4 files changed, 76 insertions(+), 16 deletions(-) diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index db6c2846f..d68d3b197 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -1,6 +1,6 @@ from math import fabs, isnan, nan -from ._rkdenseinterp import DENSE_SIG +from ._const import DENSE_SIG from ..ieee754 import EPS from ...jit import hjit @@ -132,9 +132,10 @@ def brentq_hf( return xcur, BRENTQ_CONVERR -@hjit(f"Tuple([f,f,i8])(F(f(f,{DENSE_SIG:s},f)),f,f,f,f,f,{DENSE_SIG:s},f)") +@hjit(f"Tuple([f,f,i8])(F(f(i8,f,{DENSE_SIG:s},f)),i8,f,f,f,f,f,{DENSE_SIG:s},f)") def brentq_dense_hf( func, # callback_type + idx, xa, # double xb, # double xtol, # double @@ -164,11 +165,11 @@ def brentq_dense_hf( fpre, fcur, fblk = 0.0, 0.0, 0.0 spre, scur = 0.0, 0.0 - fpre = func(xpre, sol1, sol2, sol3, sol4, sol5, argk) + fpre = func(idx, xpre, sol1, sol2, sol3, sol4, sol5, argk) if isnan(fpre): return xpre, 0.0, BRENTQ_ERROR - fcur = func(xcur, sol1, sol2, sol3, sol4, sol5, argk) + fcur = func(idx, xcur, sol1, sol2, sol3, sol4, sol5, argk) if isnan(fcur): return xcur, 0.0, BRENTQ_ERROR @@ -225,7 +226,7 @@ def brentq_dense_hf( else: xcur += delta if sbis > 0 else -delta - fcur = func(xcur, sol1, sol2, sol3, sol4, sol5, argk) + fcur = func(idx, xcur, sol1, sol2, sol3, sol4, sol5, argk) if isnan(fcur): return xcur, 0.0, BRENTQ_ERROR diff --git a/src/hapsira/core/math/ivp/_const.py b/src/hapsira/core/math/ivp/_const.py index f2bc7ea7d..69f54fedc 100644 --- a/src/hapsira/core/math/ivp/_const.py +++ b/src/hapsira/core/math/ivp/_const.py @@ -36,3 +36,5 @@ + ",".join(["Tuple([" + ",".join(["f"] * N_RV) + "])"] * INTERPOLATOR_POWER) + "])" ) + +DENSE_SIG = f"f,f,V,V,{FSIG:s}" diff --git a/src/hapsira/core/math/ivp/_rkdenseinterp.py b/src/hapsira/core/math/ivp/_rkdenseinterp.py index 2d1e7288a..790300258 100644 --- a/src/hapsira/core/math/ivp/_rkdenseinterp.py +++ b/src/hapsira/core/math/ivp/_rkdenseinterp.py @@ -1,4 +1,4 @@ -from ._const import FSIG +from ._const import DENSE_SIG from ..linalg import add_VV_hf, mul_Vs_hf from ...jit import hjit @@ -6,13 +6,9 @@ __all__ = [ "dop853_dense_interp_brentq_hb", "dop853_dense_interp_hf", - "DENSE_SIG", ] -DENSE_SIG = f"f,f,V,V,{FSIG:s}" - - @hjit(f"Tuple([V,V])(f,{DENSE_SIG:s})") def dop853_dense_interp_hf(t, t_old, h, rr_old, vv_old, F): """ diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 2e226893c..6f73fbbfb 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -1,8 +1,9 @@ from math import nan, isnan -from typing import Callable, List, Tuple +from typing import Callable, Tuple import numpy as np +from ._const import DENSE_SIG from ._brentq import brentq_dense_hf, BRENTQ_CONVERGED, BRENTQ_MAXITER from ._rkcore import ( dop853_init_hf, @@ -34,9 +35,55 @@ ] +TEMPLATE = """ +@hjit("{RESTYPE:s}(i8,{ARGTYPES:s})", cache = False) +def dispatcher_hf(idx, {ARGUMENTS:s}): +{DISPATCHER:s} + return {ERROR:s} +""" + + +def dispatcher_hb( + funcs: Tuple[Callable, ...], + argtypes: str, + restype: str, + arguments: str, + error: str = "nan", +): + """ + Workaround for https://github.com/numba/numba/issues/9420 + """ + funcs = [ + (f"func_{id(func):x}", func) for func in funcs + ] # names are not unique, ids are + globals_, locals_ = globals(), locals() # HACK https://stackoverflow.com/a/71560563 + globals_.update({name: handle for name, handle in funcs}) + + def switch(idx): + return "if" if idx == 0 else "elif" + + code = TEMPLATE.format( + DISPATCHER="\n".join( + [ + f" {switch(idx):s} idx == {idx:d}:\n return {name:s}({arguments:s})" + for idx, (name, _) in enumerate(funcs) + ] + ), # TODO tree-like dispatch, faster + ARGTYPES=argtypes, + RESTYPE=restype, + ARGUMENTS=arguments, + ERROR=error, + ) + exec(code, globals_, locals_) # pylint: disable=W0122 + globals_["dispatcher_hf"] = locals_[ + "dispatcher_hf" + ] # HACK https://stackoverflow.com/a/71560563 + return dispatcher_hf # pylint: disable=E0602 # noqa: F821 + + def _handle_events( interpolant, - event_impl_dense_hfs: List[Callable], + event_impl_dense_hf: Callable, event_last_ts: np.ndarray, event_actives: np.ndarray, event_terminals: np.ndarray, @@ -73,7 +120,7 @@ def _handle_events( assert np.any(event_actives) # nothing active - EVENTS = len(event_impl_dense_hfs) # TODO compile as const + EVENTS = len(event_last_ts) # TODO compile as const pivot = nan # set initial value terminate = False @@ -83,7 +130,8 @@ def _handle_events( continue event_last_ts[idx], root, status = brentq_dense_hf( - event_impl_dense_hfs[idx], + event_impl_dense_hf, + idx, t_old, t, 4 * EPS, @@ -163,6 +211,19 @@ def solve_ivp( EVENTS = len(event_impl_hfs) # TODO compile as const + event_impl_hf = dispatcher_hb( + funcs=event_impl_hfs, + argtypes="f,V,V,f", + restype="f", + arguments="t, rr, vv, k", + ) + event_impl_dense_hf = dispatcher_hb( + funcs=event_impl_dense_hfs, + argtypes=f"f,{DENSE_SIG:s},f", + restype="f", + arguments="t, t_old, h, rr_old, vv_old, F, argk", + ) + solver = dop853_init_hf(fun, t0, rr, vv, tf, argk, rtol, atol) ts = [t0] interpolants = [] @@ -182,7 +243,7 @@ def solve_ivp( """ for idx in range(EVENTS): - event_g_olds[idx] = event_impl_hfs[idx](t0, rr, vv, argk) + event_g_olds[idx] = event_impl_hf(idx, t0, rr, vv, argk) event_last_ts[idx] = t0 status = None @@ -230,7 +291,7 @@ def solve_ivp( if at_least_one_active: root, terminate = _handle_events( interpolant, - event_impl_dense_hfs, # TODO + event_impl_dense_hf, # TODO event_last_ts, event_actives, # TODO event_terminals, # TODO From 7cc7707cd04625cd07d4e3c9964dcecc089db2e3 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 4 Feb 2024 13:28:17 +0100 Subject: [PATCH 276/346] mv event handling into solver --- src/hapsira/core/math/ivp/_solve.py | 134 +++++++++------------------- 1 file changed, 40 insertions(+), 94 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 6f73fbbfb..dd4188210 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -81,89 +81,6 @@ def switch(idx): return dispatcher_hf # pylint: disable=E0602 # noqa: F821 -def _handle_events( - interpolant, - event_impl_dense_hf: Callable, - event_last_ts: np.ndarray, - event_actives: np.ndarray, - event_terminals: np.ndarray, - t_old: float, - t: float, - argk: float, -): - """Helper function to handle events. - - Parameters - ---------- - sol : DenseOutput - Function ``sol(t)`` which evaluates an ODE solution between `t_old` - and `t`. - events : list of callables, length n_events - Event functions with signatures ``event(t, y)``. - active_events : ndarray - Indices of events which occurred. - terminals : ndarray, shape (n_events,) - Which events are terminal. - t_old, t : float - Previous and new values of time. - - Returns - ------- - root_indices : ndarray - Indices of events which take zero between `t_old` and `t` and before - a possible termination. - roots : ndarray - Values of t at which events occurred. - terminate : bool - Whether a terminal event occurred. - """ - - assert np.any(event_actives) # nothing active - - EVENTS = len(event_last_ts) # TODO compile as const - - pivot = nan # set initial value - terminate = False - - for idx in range(EVENTS): - if not event_actives[idx]: - continue - - event_last_ts[idx], root, status = brentq_dense_hf( - event_impl_dense_hf, - idx, - t_old, - t, - 4 * EPS, - 4 * EPS, - BRENTQ_MAXITER, - *interpolant, - argk, - ) - assert status == BRENTQ_CONVERGED - - if event_terminals[idx]: - terminate = True - - if isnan(pivot): - pivot = root - continue - - if t > t_old: # smallest root of all active events - if root < pivot: - pivot = root - continue - - # largest root of all active events - if root > pivot: - pivot = root - raise ValueError("not t > t_old", t, t_old) # TODO remove - - assert not isnan(pivot) - - return pivot if terminate else nan, terminate - - @hjit("b1(f,f,f)") def _event_is_active_hf(g_old, g_new, direction): """Find which event occurred during an integration step. @@ -289,19 +206,48 @@ def solve_ivp( at_least_one_active = True if at_least_one_active: - root, terminate = _handle_events( - interpolant, - event_impl_dense_hf, # TODO - event_last_ts, - event_actives, # TODO - event_terminals, # TODO - t_old, - t, - argk, - ) + root_pivot = nan # set initial value + terminate = False + + for idx in range(EVENTS): + if not event_actives[idx]: + continue + + event_last_ts[idx], root, status = brentq_dense_hf( + event_impl_dense_hf, + idx, + t_old, + t, + 4 * EPS, + 4 * EPS, + BRENTQ_MAXITER, + *interpolant, + argk, + ) + assert status == BRENTQ_CONVERGED + + if event_terminals[idx]: + terminate = True + + if isnan(root_pivot): + root_pivot = root + continue + + if t > t_old: # smallest root of all active events + if root < root_pivot: + root_pivot = root + continue + + # largest root of all active events + if root > root_pivot: + root_pivot = root + raise ValueError("not t > t_old", t, t_old) # TODO remove + + assert not isnan(root_pivot) + if terminate: status = 1 - t = root + t = root_pivot for idx in range(EVENTS): event_g_olds[idx] = event_g_news[idx] From ab9aaa614ae4a4ae714af6d03e1f90348b08faf6 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 4 Feb 2024 13:47:11 +0100 Subject: [PATCH 277/346] fix type --- src/hapsira/core/math/ivp/_solve.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index dd4188210..571b14ce8 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -49,7 +49,7 @@ def dispatcher_hb( restype: str, arguments: str, error: str = "nan", -): +) -> Callable: """ Workaround for https://github.com/numba/numba/issues/9420 """ From 253b008a2ec02367a5fd2f93a3e75b45dc0486fe Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 4 Feb 2024 20:48:59 +0100 Subject: [PATCH 278/346] rm last allocs; mv all of event handling into solver --- src/hapsira/core/math/ivp/_solve.py | 116 +++++++++++-------------- src/hapsira/core/propagation/cowell.py | 37 +++----- 2 files changed, 63 insertions(+), 90 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 571b14ce8..393e3b647 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -35,6 +35,9 @@ ] +SOLVE_TERMINATED = 1 + + TEMPLATE = """ @hjit("{RESTYPE:s}(i8,{ARGTYPES:s})", cache = False) def dispatcher_hf(idx, {ARGUMENTS:s}): @@ -106,10 +109,11 @@ def _event_is_active_hf(g_old, g_new, direction): def solve_ivp( fun: Callable, - t0: float, - tf: float, + tofs: np.ndarray, rr: Tuple[float, float, float], vv: Tuple[float, float, float], + rrs: np.ndarray, + vvs: np.ndarray, argk: float, rtol: float, atol: float, @@ -127,6 +131,7 @@ def solve_ivp( """ EVENTS = len(event_impl_hfs) # TODO compile as const + T0 = 0.0 event_impl_hf = dispatcher_hb( funcs=event_impl_hfs, @@ -141,27 +146,14 @@ def solve_ivp( arguments="t, t_old, h, rr_old, vv_old, F, argk", ) - solver = dop853_init_hf(fun, t0, rr, vv, tf, argk, rtol, atol) - ts = [t0] - interpolants = [] - - _ = """ - Event: - - impl_hf (callable) -> compiled tuple - - impl_dense_hf (callable) -> compiled tuple - - terminal (const) -> input array - - direction (const) -> input array - - is_active - - g_old - - g_new - - last_t -> output array - N events -> compiled const int? - for-loop??? - """ + solver = dop853_init_hf(fun, T0, rr, vv, tofs[-1], argk, rtol, atol) + + t_idx = 0 + t_last = T0 - for idx in range(EVENTS): - event_g_olds[idx] = event_impl_hf(idx, t0, rr, vv, argk) - event_last_ts[idx] = t0 + for event_idx in range(EVENTS): + event_g_olds[event_idx] = event_impl_hf(event_idx, T0, rr, vv, argk) + event_last_ts[event_idx] = T0 status = None while status is None: @@ -192,30 +184,35 @@ def solve_ivp( ) at_least_one_active = False - for idx in range(EVENTS): - event_g_news[idx] = event_impl_hfs[idx]( + for event_idx in range(EVENTS): + event_g_news[event_idx] = event_impl_hfs[event_idx]( t, solver[DOP853_RR], solver[DOP853_VV], argk ) - event_last_ts[idx] = t - event_actives[idx] = _event_is_active_hf( - event_g_olds[idx], - event_g_news[idx], - event_directions[idx], + event_last_ts[event_idx] = t + event_actives[event_idx] = _event_is_active_hf( + event_g_olds[event_idx], + event_g_news[event_idx], + event_directions[event_idx], ) - if event_actives[idx]: + if event_actives[event_idx]: at_least_one_active = True if at_least_one_active: root_pivot = nan # set initial value terminate = False - for idx in range(EVENTS): - if not event_actives[idx]: + for event_idx in range(EVENTS): + if not event_actives[event_idx]: continue - event_last_ts[idx], root, status = brentq_dense_hf( + if not event_terminals[event_idx]: + continue + + terminate = True + + event_last_ts[event_idx], root, status = brentq_dense_hf( event_impl_dense_hf, - idx, + event_idx, t_old, t, 4 * EPS, @@ -226,9 +223,6 @@ def solve_ivp( ) assert status == BRENTQ_CONVERGED - if event_terminals[idx]: - terminate = True - if isnan(root_pivot): root_pivot = root continue @@ -243,36 +237,30 @@ def solve_ivp( root_pivot = root raise ValueError("not t > t_old", t, t_old) # TODO remove - assert not isnan(root_pivot) - if terminate: - status = 1 + assert not isnan(root_pivot) + status = SOLVE_TERMINATED t = root_pivot - for idx in range(EVENTS): - event_g_olds[idx] = event_g_news[idx] + for event_idx in range(EVENTS): + event_g_olds[event_idx] = event_g_news[event_idx] - if not ts[-1] <= t: - raise ValueError("not ts[-1] <= t", ts[-1], t) - interpolants.append(interpolant) - ts.append(t) + if not t_last <= t: + raise ValueError("not t_last <= t", t_last, t) - assert len(ts) >= 2 - assert len(ts) == len(interpolants) + 1 - assert ( - (len(ts) == 2 and ts[0] == ts[1]) - or all(a - b > 0 for a, b in zip(ts[:-1], ts[1:])) - or all(b - a > 0 for a, b in zip(ts[:-1], ts[1:])) - ) + while t_idx < tofs.shape[0] and tofs[t_idx] < t: + rrs[t_idx, :], vvs[t_idx, :] = dop853_dense_interp_hf( + tofs[t_idx], *interpolant + ) + t_idx += 1 + if status == SOLVE_TERMINATED or tofs[t_idx] == t: + rrs[t_idx, :], vvs[t_idx, :] = dop853_dense_interp_hf(t, *interpolant) + t_idx += 1 + + t_last = t + + # while t_idx < tofs.shape[0] and status != SOLVE_TERMINATED: # fill up if not terminated + # rrs[t_idx, :], vvs[t_idx, :] = dop853_dense_interp_hf(tofs[t_idx], *interpolant) + # t_idx += 1 - def ode_solution( - t: float, - ) -> Tuple[Tuple[float, float, float], Tuple[float, float, float]]: - """ - Evaluate the solution - """ - idx = np.searchsorted(ts, t, side="left") - segment = min(max(idx - 1, 0), len(interpolants) - 1) - return dop853_dense_interp_hf(t, *interpolants[segment]) - - return ode_solution, status >= 0 + return t_idx, status >= 0 diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index d62a73694..cc0ad2cf7 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -33,6 +33,9 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody EVENTS = len(events) # TODO compile as const + rrs = np.full((len(tofs), 3), fill_value=np.nan, dtype=float) + vvs = np.full((len(tofs), 3), fill_value=np.nan, dtype=float) + event_impl_hfs = tuple( event.impl_hf for event in events ) # TODO compile into kernel @@ -58,12 +61,13 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody (EVENTS,), fill_value=np.nan, dtype=float ) # gufunc param TODO reset to nan - sol, success = solve_ivp( - f, - 0.0, - float(max(tofs)), - array_to_V_hf(r), - array_to_V_hf(v), + length, success = solve_ivp( + fun=f, + tofs=tofs, + rr=array_to_V_hf(r), + vv=array_to_V_hf(v), + rrs=rrs, + vvs=vvs, argk=k, rtol=rtol, atol=atol, @@ -82,23 +86,4 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody for idx in range(EVENTS): events[idx].last_t_raw = event_last_ts[idx] - if len(events) > 0: - # Collect only the terminal events - terminal_events = [event for event in events if event.terminal] - - # If there are no terminal events, then the last time of integration is the - # greatest one from the original array of propagation times - if len(terminal_events) > 0: - # Filter the event which triggered first - last_t = min(event.last_t_raw for event in terminal_events) - tofs = [tof for tof in tofs if tof < last_t] - tofs.append(last_t) - - rrs = [] - vvs = [] - for t in tofs: - r_new, v_new = sol(t) - rrs.append(r_new) - vvs.append(v_new) - - return rrs, vvs + return rrs[:length, :], vvs[:length, :] From ed011c98b395575d978a9f6b0b0c24f7828c2070 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 4 Feb 2024 20:49:37 +0100 Subject: [PATCH 279/346] cleanup --- src/hapsira/core/math/ivp/_solve.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 393e3b647..de2de887a 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -259,8 +259,4 @@ def solve_ivp( t_last = t - # while t_idx < tofs.shape[0] and status != SOLVE_TERMINATED: # fill up if not terminated - # rrs[t_idx, :], vvs[t_idx, :] = dop853_dense_interp_hf(tofs[t_idx], *interpolant) - # t_idx += 1 - return t_idx, status >= 0 From ba10b2ffc7ac22542abbb646cda2d414cb37f643 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 4 Feb 2024 21:02:07 +0100 Subject: [PATCH 280/346] fix error handling --- src/hapsira/core/math/ivp/_solve.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index de2de887a..f56349e54 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -35,9 +35,11 @@ ] +SOLVE_RUNNING = -2 +SOLVE_FAILED = -1 +SOLVE_FINISHED = 0 SOLVE_TERMINATED = 1 - TEMPLATE = """ @hjit("{RESTYPE:s}(i8,{ARGTYPES:s})", cache = False) def dispatcher_hf(idx, {ARGUMENTS:s}): @@ -155,14 +157,14 @@ def solve_ivp( event_g_olds[event_idx] = event_impl_hf(event_idx, T0, rr, vv, argk) event_last_ts[event_idx] = T0 - status = None - while status is None: + status = SOLVE_RUNNING + while status == SOLVE_RUNNING: solver = dop853_step_hf(*solver) if solver[DOP853_STATUS] == DOP853_FINISHED: - status = 0 + status = SOLVE_FINISHED elif solver[DOP853_STATUS] == DOP853_FAILED: - status = -1 + status = SOLVE_FAILED break t_old = solver[DOP853_T_OLD] @@ -210,7 +212,7 @@ def solve_ivp( terminate = True - event_last_ts[event_idx], root, status = brentq_dense_hf( + event_last_ts[event_idx], root, brentq_status = brentq_dense_hf( event_impl_dense_hf, event_idx, t_old, @@ -221,7 +223,8 @@ def solve_ivp( *interpolant, argk, ) - assert status == BRENTQ_CONVERGED + if brentq_status != BRENTQ_CONVERGED: + return t_idx, False # failed on event if isnan(root_pivot): root_pivot = root From 65430db5bfeb2911642fbf1bf22e3ecb08276f00 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 4 Feb 2024 22:25:14 +0100 Subject: [PATCH 281/346] fix names --- src/hapsira/core/math/ivp/_solve.py | 4 ++-- src/hapsira/core/propagation/cowell.py | 12 +++++++----- src/hapsira/twobody/propagation/cowell.py | 12 ++++++------ 3 files changed, 15 insertions(+), 13 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index f56349e54..5dc830d85 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -110,7 +110,7 @@ def _event_is_active_hf(g_old, g_new, direction): def solve_ivp( - fun: Callable, + func: Callable, tofs: np.ndarray, rr: Tuple[float, float, float], vv: Tuple[float, float, float], @@ -148,7 +148,7 @@ def solve_ivp( arguments="t, t_old, h, rr_old, vv_old, F, argk", ) - solver = dop853_init_hf(fun, T0, rr, vv, tofs[-1], argk, rtol, atol) + solver = dop853_init_hf(func, T0, rr, vv, tofs[-1], argk, rtol, atol) t_idx = 0 t_last = T0 diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index cc0ad2cf7..4e1503c2f 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -6,11 +6,13 @@ __all__ = [ - "cowell", + "cowell_vb", ] -def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody_hf): +def cowell_vb( + k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), func=func_twobody_hf +): """ Scalar cowell @@ -21,12 +23,12 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody rtol : float atol : float events : Optional[List[Event]] - f : Callable + func : Callable Can be reversed: https://github.com/poliastro/poliastro/issues/1630 """ - assert hasattr(f, "djit") # DEBUG check for compiler flag + assert hasattr(func, "djit") # DEBUG check for compiler flag assert isinstance(rtol, float) assert all(tof >= 0 for tof in tofs) assert sorted(tofs) == list(tofs) @@ -62,7 +64,7 @@ def cowell(k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody ) # gufunc param TODO reset to nan length, success = solve_ivp( - fun=f, + func=func, tofs=tofs, rr=array_to_V_hf(r), vv=array_to_V_hf(v), diff --git a/src/hapsira/twobody/propagation/cowell.py b/src/hapsira/twobody/propagation/cowell.py index e7410e125..b002c78c1 100644 --- a/src/hapsira/twobody/propagation/cowell.py +++ b/src/hapsira/twobody/propagation/cowell.py @@ -2,7 +2,7 @@ from astropy import units as u -from hapsira.core.propagation.cowell import cowell +from hapsira.core.propagation.cowell import cowell_vb from hapsira.core.propagation.base import func_twobody_hf from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import RVState @@ -30,19 +30,19 @@ class CowellPropagator: def __init__(self, rtol=1e-11, events=tuple(), f=func_twobody_hf): self._rtol = rtol self._events = events - self._f = f + self._func = f def propagate(self, state, tof): state = state.to_vectors() tofs = tof.reshape(-1) - rrs, vvs = cowell( + rrs, vvs = cowell_vb( state.attractor.k.to_value(u.km**3 / u.s**2), *state.to_value(), tofs.to_value(u.s), self._rtol, events=self._events, - f=self._f, + func=self._func, ) r = rrs[-1] << u.km v = vvs[-1] << (u.km / u.s) @@ -53,13 +53,13 @@ def propagate(self, state, tof): def propagate_many(self, state, tofs): state = state.to_vectors() - rrs, vvs = cowell( + rrs, vvs = cowell_vb( state.attractor.k.to_value(u.km**3 / u.s**2), *state.to_value(), tofs.to_value(u.s), self._rtol, events=self._events, - f=self._f, + func=self._func, ) # TODO: This should probably return a RVStateArray instead, From dc4cd82330d27b1b0e30d69256ebfd0e357682a5 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 5 Feb 2024 11:27:58 +0100 Subject: [PATCH 282/346] jit cowell --- src/hapsira/core/math/ivp/__init__.py | 48 +++- src/hapsira/core/math/ivp/_solve.py | 196 +-------------- src/hapsira/core/propagation/cowell.py | 287 ++++++++++++++++------ src/hapsira/twobody/propagation/cowell.py | 88 +++++-- 4 files changed, 333 insertions(+), 286 deletions(-) diff --git a/src/hapsira/core/math/ivp/__init__.py b/src/hapsira/core/math/ivp/__init__.py index 173e3926d..7d1562754 100644 --- a/src/hapsira/core/math/ivp/__init__.py +++ b/src/hapsira/core/math/ivp/__init__.py @@ -7,9 +7,32 @@ BRENTQ_RTOL, BRENTQ_MAXITER, brentq_hf, + brentq_dense_hf, +) +from ._const import DENSE_SIG +from ._rkcore import ( + dop853_init_hf, + dop853_step_hf, + DOP853_FINISHED, + DOP853_FAILED, + DOP853_ARGK, + DOP853_FR, + DOP853_FUN, + DOP853_FV, + DOP853_H_PREVIOUS, + DOP853_K, + DOP853_RR, + DOP853_RR_OLD, + DOP853_STATUS, + DOP853_T, + DOP853_T_OLD, + DOP853_VV, + DOP853_VV_OLD, ) from ._rkdenseinterp import dop853_dense_interp_brentq_hb, dop853_dense_interp_hf -from ._solve import solve_ivp +from ._rkdenseoutput import dop853_dense_output_hf +from ._solve import event_is_active_hf, dispatcher_hb + __all__ = [ "BRENTQ_CONVERGED", @@ -19,8 +42,29 @@ "BRENTQ_XTOL", "BRENTQ_RTOL", "BRENTQ_MAXITER", + "DENSE_SIG", + "DOP853_FINISHED", + "DOP853_FAILED", + "DOP853_ARGK", + "DOP853_FR", + "DOP853_FUN", + "DOP853_FV", + "DOP853_H_PREVIOUS", + "DOP853_K", + "DOP853_RR", + "DOP853_RR_OLD", + "DOP853_STATUS", + "DOP853_T", + "DOP853_T_OLD", + "DOP853_VV", + "DOP853_VV_OLD", "brentq_hf", + "brentq_dense_hf", + "dispatcher_hb", "dop853_dense_interp_brentq_hb", "dop853_dense_interp_hf", - "solve_ivp", + "dop853_dense_output_hf", + "dop853_init_hf", + "dop853_step_hf", + "event_is_active_hf", ] diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index 5dc830d85..ad6baf421 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -1,45 +1,15 @@ -from math import nan, isnan +from math import nan from typing import Callable, Tuple -import numpy as np - -from ._const import DENSE_SIG -from ._brentq import brentq_dense_hf, BRENTQ_CONVERGED, BRENTQ_MAXITER -from ._rkcore import ( - dop853_init_hf, - dop853_step_hf, - DOP853_FINISHED, - DOP853_FAILED, - DOP853_ARGK, - DOP853_FR, - DOP853_FUN, - DOP853_FV, - DOP853_H_PREVIOUS, - DOP853_K, - DOP853_RR, - DOP853_RR_OLD, - DOP853_STATUS, - DOP853_T, - DOP853_T_OLD, - DOP853_VV, - DOP853_VV_OLD, -) -from ._rkdenseinterp import dop853_dense_interp_hf -from ._rkdenseoutput import dop853_dense_output_hf -from ..ieee754 import EPS from ...jit import hjit __all__ = [ - "solve_ivp", + "event_is_active_hf", + "dispatcher_hb", ] -SOLVE_RUNNING = -2 -SOLVE_FAILED = -1 -SOLVE_FINISHED = 0 -SOLVE_TERMINATED = 1 - TEMPLATE = """ @hjit("{RESTYPE:s}(i8,{ARGTYPES:s})", cache = False) def dispatcher_hf(idx, {ARGUMENTS:s}): @@ -47,6 +17,8 @@ def dispatcher_hf(idx, {ARGUMENTS:s}): return {ERROR:s} """ +_ = nan # keep import alive + def dispatcher_hb( funcs: Tuple[Callable, ...], @@ -87,7 +59,7 @@ def switch(idx): @hjit("b1(f,f,f)") -def _event_is_active_hf(g_old, g_new, direction): +def event_is_active_hf(g_old, g_new, direction): """Find which event occurred during an integration step. Parameters @@ -107,159 +79,3 @@ def _event_is_active_hf(g_old, g_new, direction): either = up | down active = up & (direction > 0) | down & (direction < 0) | either & (direction == 0) return active - - -def solve_ivp( - func: Callable, - tofs: np.ndarray, - rr: Tuple[float, float, float], - vv: Tuple[float, float, float], - rrs: np.ndarray, - vvs: np.ndarray, - argk: float, - rtol: float, - atol: float, - event_impl_hfs: Tuple[Callable, ...], - event_impl_dense_hfs: Tuple[Callable, ...], - event_terminals: np.ndarray, - event_directions: np.ndarray, - event_actives: np.ndarray, - event_g_olds: np.ndarray, - event_g_news: np.ndarray, - event_last_ts: np.ndarray, -) -> Tuple[Callable, bool]: - """ - Solve an initial value problem for a system of ODEs. - """ - - EVENTS = len(event_impl_hfs) # TODO compile as const - T0 = 0.0 - - event_impl_hf = dispatcher_hb( - funcs=event_impl_hfs, - argtypes="f,V,V,f", - restype="f", - arguments="t, rr, vv, k", - ) - event_impl_dense_hf = dispatcher_hb( - funcs=event_impl_dense_hfs, - argtypes=f"f,{DENSE_SIG:s},f", - restype="f", - arguments="t, t_old, h, rr_old, vv_old, F, argk", - ) - - solver = dop853_init_hf(func, T0, rr, vv, tofs[-1], argk, rtol, atol) - - t_idx = 0 - t_last = T0 - - for event_idx in range(EVENTS): - event_g_olds[event_idx] = event_impl_hf(event_idx, T0, rr, vv, argk) - event_last_ts[event_idx] = T0 - - status = SOLVE_RUNNING - while status == SOLVE_RUNNING: - solver = dop853_step_hf(*solver) - - if solver[DOP853_STATUS] == DOP853_FINISHED: - status = SOLVE_FINISHED - elif solver[DOP853_STATUS] == DOP853_FAILED: - status = SOLVE_FAILED - break - - t_old = solver[DOP853_T_OLD] - t = solver[DOP853_T] - - interpolant = dop853_dense_output_hf( - solver[DOP853_FUN], - solver[DOP853_ARGK], - solver[DOP853_T_OLD], - solver[DOP853_T], - solver[DOP853_H_PREVIOUS], - solver[DOP853_RR], - solver[DOP853_VV], - solver[DOP853_RR_OLD], - solver[DOP853_VV_OLD], - solver[DOP853_FR], - solver[DOP853_FV], - solver[DOP853_K], - ) - - at_least_one_active = False - for event_idx in range(EVENTS): - event_g_news[event_idx] = event_impl_hfs[event_idx]( - t, solver[DOP853_RR], solver[DOP853_VV], argk - ) - event_last_ts[event_idx] = t - event_actives[event_idx] = _event_is_active_hf( - event_g_olds[event_idx], - event_g_news[event_idx], - event_directions[event_idx], - ) - if event_actives[event_idx]: - at_least_one_active = True - - if at_least_one_active: - root_pivot = nan # set initial value - terminate = False - - for event_idx in range(EVENTS): - if not event_actives[event_idx]: - continue - - if not event_terminals[event_idx]: - continue - - terminate = True - - event_last_ts[event_idx], root, brentq_status = brentq_dense_hf( - event_impl_dense_hf, - event_idx, - t_old, - t, - 4 * EPS, - 4 * EPS, - BRENTQ_MAXITER, - *interpolant, - argk, - ) - if brentq_status != BRENTQ_CONVERGED: - return t_idx, False # failed on event - - if isnan(root_pivot): - root_pivot = root - continue - - if t > t_old: # smallest root of all active events - if root < root_pivot: - root_pivot = root - continue - - # largest root of all active events - if root > root_pivot: - root_pivot = root - raise ValueError("not t > t_old", t, t_old) # TODO remove - - if terminate: - assert not isnan(root_pivot) - status = SOLVE_TERMINATED - t = root_pivot - - for event_idx in range(EVENTS): - event_g_olds[event_idx] = event_g_news[event_idx] - - if not t_last <= t: - raise ValueError("not t_last <= t", t_last, t) - - while t_idx < tofs.shape[0] and tofs[t_idx] < t: - rrs[t_idx, :], vvs[t_idx, :] = dop853_dense_interp_hf( - tofs[t_idx], *interpolant - ) - t_idx += 1 - if status == SOLVE_TERMINATED or tofs[t_idx] == t: - rrs[t_idx, :], vvs[t_idx, :] = dop853_dense_interp_hf(t, *interpolant) - t_idx += 1 - - t_last = t - - return t_idx, status >= 0 diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 4e1503c2f..b18b26aeb 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -1,91 +1,236 @@ -import numpy as np +from math import isnan, nan +from typing import Callable, Tuple -from ..jit import array_to_V_hf -from ..math.ivp import solve_ivp +from numpy import ndarray + +from ..jit import gjit, array_to_V_hf +from ..math.ieee754 import EPS +from ..math.ivp import ( + BRENTQ_CONVERGED, + BRENTQ_MAXITER, + DENSE_SIG, + DOP853_FINISHED, + DOP853_FAILED, + DOP853_ARGK, + DOP853_FR, + DOP853_FUN, + DOP853_FV, + DOP853_H_PREVIOUS, + DOP853_K, + DOP853_RR, + DOP853_RR_OLD, + DOP853_STATUS, + DOP853_T, + DOP853_T_OLD, + DOP853_VV, + DOP853_VV_OLD, + brentq_dense_hf, + event_is_active_hf, + dispatcher_hb, + dop853_dense_interp_hf, + dop853_dense_output_hf, + dop853_init_hf, + dop853_step_hf, +) from ..propagation.base import func_twobody_hf __all__ = [ "cowell_vb", + "SOLVE_BRENTQFAILED", + "SOLVE_FAILED", + "SOLVE_RUNNING", + "SOLVE_FINISHED", + "SOLVE_TERMINATED", ] +SOLVE_BRENTQFAILED = -3 +SOLVE_FAILED = -2 +SOLVE_RUNNING = -1 +SOLVE_FINISHED = 0 +SOLVE_TERMINATED = 1 + + def cowell_vb( - k, r, v, tofs, rtol=1e-11, atol=1e-12, events=tuple(), func=func_twobody_hf -): + events: Tuple = tuple(), + func: Callable = func_twobody_hf, +) -> Callable: """ - Scalar cowell - - k : float - r : ndarray (3,) - v : ndarray (3,) - tofs : array of relative times [seconds] - rtol : float - atol : float - events : Optional[List[Event]] - func : Callable - - Can be reversed: https://github.com/poliastro/poliastro/issues/1630 + Builds vectorized cowell """ assert hasattr(func, "djit") # DEBUG check for compiler flag - assert isinstance(rtol, float) - assert all(tof >= 0 for tof in tofs) - assert sorted(tofs) == list(tofs) EVENTS = len(events) # TODO compile as const - rrs = np.full((len(tofs), 3), fill_value=np.nan, dtype=float) - vvs = np.full((len(tofs), 3), fill_value=np.nan, dtype=float) - - event_impl_hfs = tuple( - event.impl_hf for event in events - ) # TODO compile into kernel - event_impl_dense_hfs = tuple( - event.impl_dense_hf for event in events - ) # TODO compile into kernel - event_terminals = np.array( - [event.terminal for event in events], dtype=bool - ) # gufunc param static - event_directions = np.array( - [event.direction for event in events], dtype=float - ) # gufunc param static - event_actives = np.full( - (EVENTS,), fill_value=np.nan, dtype=bool - ) # gufunc param TODO reset to nan - event_g_olds = np.full( - (EVENTS,), fill_value=np.nan, dtype=float - ) # gufunc param TODO reset to nan - event_g_news = np.full( - (EVENTS,), fill_value=np.nan, dtype=float - ) # gufunc param TODO reset to nan - event_last_ts = np.full( - (EVENTS,), fill_value=np.nan, dtype=float - ) # gufunc param TODO reset to nan - - length, success = solve_ivp( - func=func, - tofs=tofs, - rr=array_to_V_hf(r), - vv=array_to_V_hf(v), - rrs=rrs, - vvs=vvs, - argk=k, - rtol=rtol, - atol=atol, - event_impl_hfs=event_impl_hfs, - event_impl_dense_hfs=event_impl_dense_hfs, - event_terminals=event_terminals, - event_directions=event_directions, - event_actives=event_actives, - event_g_olds=event_g_olds, - event_g_news=event_g_news, - event_last_ts=event_last_ts, + event_impl_hf = dispatcher_hb( + funcs=tuple(event.impl_hf for event in events), + argtypes="f,V,V,f", + restype="f", + arguments="t, rr, vv, k", + ) + event_impl_dense_hf = dispatcher_hb( + funcs=tuple(event.impl_dense_hf for event in events), + argtypes=f"f,{DENSE_SIG:s},f", + restype="f", + arguments="t, t_old, h, rr_old, vv_old, F, argk", ) - if not success: - raise RuntimeError("Integration failed") - for idx in range(EVENTS): - events[idx].last_t_raw = event_last_ts[idx] + @gjit( + "void(f[:],f[:],f[:],f,f,f,b1[:],f[:],f[:],f[:],f[:],f[:],i8[:],i8[:],f[:,:],f[:,:])", + "(n),(m),(m),(),(),(),(o),(o)->(o),(o),(o),(o),(),(),(n,m),(n,m)", + cache=False, + ) # n: tofs, m: dims, o: events + def cowell_gf( + tofs: ndarray, + rr: Tuple[float, float, float], + vv: Tuple[float, float, float], + argk: float, + rtol: float, + atol: float, + event_terminals: ndarray, + event_directions: ndarray, + event_g_olds: ndarray, # (out) + event_g_news: ndarray, # (out) + event_actives: ndarray, # out + event_last_ts: ndarray, # out + status: int, # out + t_idx: int, # out + rrs: ndarray, # out + vvs: ndarray, # out + ): # -> void(..., rrs, vvs, success) + """ + Solve an initial value problem for a system of ODEs. + + Can theoretically be reversed: https://github.com/poliastro/poliastro/issues/1630 + """ + + # assert isinstance(rtol, float) + # assert all(tof >= 0 for tof in tofs) + # assert sorted(tofs) == list(tofs) + + T0 = 0.0 + + solver = dop853_init_hf( + func, T0, array_to_V_hf(rr), array_to_V_hf(vv), tofs[-1], argk, rtol, atol + ) + + t_idx[0] = 0 + t_last = T0 + + for event_idx in range(EVENTS): + event_g_olds[event_idx] = event_impl_hf( + event_idx, T0, array_to_V_hf(rr), array_to_V_hf(vv), argk + ) + event_last_ts[event_idx] = T0 + + status[0] = SOLVE_RUNNING + while status[0] == SOLVE_RUNNING: + solver = dop853_step_hf(*solver) + + if solver[DOP853_STATUS] == DOP853_FINISHED: + status[0] = SOLVE_FINISHED + elif solver[DOP853_STATUS] == DOP853_FAILED: + status[0] = SOLVE_FAILED + break + + t_old = solver[DOP853_T_OLD] + t = solver[DOP853_T] + + interpolant = dop853_dense_output_hf( + solver[DOP853_FUN], + solver[DOP853_ARGK], + solver[DOP853_T_OLD], + solver[DOP853_T], + solver[DOP853_H_PREVIOUS], + solver[DOP853_RR], + solver[DOP853_VV], + solver[DOP853_RR_OLD], + solver[DOP853_VV_OLD], + solver[DOP853_FR], + solver[DOP853_FV], + solver[DOP853_K], + ) + + at_least_one_active = False + for event_idx in range(EVENTS): + event_g_news[event_idx] = event_impl_hf( + event_idx, t, solver[DOP853_RR], solver[DOP853_VV], argk + ) + event_last_ts[event_idx] = t + event_actives[event_idx] = event_is_active_hf( + event_g_olds[event_idx], + event_g_news[event_idx], + event_directions[event_idx], + ) + if event_actives[event_idx]: + at_least_one_active = True + + if at_least_one_active: + root_pivot = nan # set initial value + terminate = False + + for event_idx in range(EVENTS): + if not event_actives[event_idx]: + continue + + if not event_terminals[event_idx]: + continue + + terminate = True + + event_last_ts[event_idx], root, brentq_status = brentq_dense_hf( + event_impl_dense_hf, + event_idx, + t_old, + t, + 4 * EPS, + 4 * EPS, + BRENTQ_MAXITER, + *interpolant, + argk, + ) + if brentq_status != BRENTQ_CONVERGED: + status[0] = SOLVE_BRENTQFAILED + return # failed on event + + if isnan(root_pivot): + root_pivot = root + continue + + if t > t_old: # smallest root of all active events + if root < root_pivot: + root_pivot = root + continue + + # largest root of all active events + if root > root_pivot: + root_pivot = root + raise ValueError("not t > t_old", t, t_old) # TODO remove + + if terminate: + assert not isnan(root_pivot) + status[0] = SOLVE_TERMINATED + t = root_pivot + + for event_idx in range(EVENTS): + event_g_olds[event_idx] = event_g_news[event_idx] + + if not t_last <= t: + raise ValueError("not t_last <= t", t_last, t) + + while t_idx[0] < tofs.shape[0] and tofs[t_idx[0]] < t: + rrs[t_idx[0], :], vvs[t_idx[0], :] = dop853_dense_interp_hf( + tofs[t_idx[0]], *interpolant + ) + t_idx[0] += 1 + if status[0] == SOLVE_TERMINATED or tofs[t_idx[0]] == t: + rrs[t_idx[0], :], vvs[t_idx[0], :] = dop853_dense_interp_hf( + t, *interpolant + ) + t_idx[0] += 1 + + t_last = t - return rrs[:length, :], vvs[:length, :] + return cowell_gf diff --git a/src/hapsira/twobody/propagation/cowell.py b/src/hapsira/twobody/propagation/cowell.py index b002c78c1..742d179fd 100644 --- a/src/hapsira/twobody/propagation/cowell.py +++ b/src/hapsira/twobody/propagation/cowell.py @@ -1,8 +1,10 @@ import sys from astropy import units as u +import numpy as np -from hapsira.core.propagation.cowell import cowell_vb +from hapsira.core.math.ieee754 import float_ +from hapsira.core.propagation.cowell import cowell_vb, SOLVE_FINISHED, SOLVE_TERMINATED from hapsira.core.propagation.base import func_twobody_hf from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import RVState @@ -27,44 +29,84 @@ class CowellPropagator: PropagatorKind.ELLIPTIC | PropagatorKind.PARABOLIC | PropagatorKind.HYPERBOLIC ) - def __init__(self, rtol=1e-11, events=tuple(), f=func_twobody_hf): + def __init__(self, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody_hf): self._rtol = rtol + self._atol = atol self._events = events - self._func = f + self._terminals = np.array([event.terminal for event in events], dtype=bool) + self._directions = np.array([event.direction for event in events], dtype=float_) + self._cowell_gf = cowell_vb(events=events, func=f) def propagate(self, state, tof): state = state.to_vectors() tofs = tof.reshape(-1) - - rrs, vvs = cowell_vb( - state.attractor.k.to_value(u.km**3 / u.s**2), - *state.to_value(), - tofs.to_value(u.s), - self._rtol, - events=self._events, - func=self._func, + # TODO make sure tofs is sorted + + r0, v0 = state.to_value() + ( # pylint: disable=E0633,E1120 + _, + _, + _, + last_ts, + status, + t_idx, + rrs, + vvs, + ) = self._cowell_gf( # pylint: disable=E0633,E1120 + tofs.to_value(u.s), # tofs + r0, # rr + v0, # vv + state.attractor.k.to_value(u.km**3 / u.s**2), # argk + self._rtol, # rtol + self._atol, # atol + self._terminals, # event_terminals + self._directions, # event_directions ) - r = rrs[-1] << u.km - v = vvs[-1] << (u.km / u.s) + + assert np.all((status == SOLVE_FINISHED) | (status == SOLVE_TERMINATED)) + + for last_t, event in zip(last_ts, self._events): + event.last_t_raw = last_t + + r = rrs[t_idx - 1] << u.km + v = vvs[t_idx - 1] << (u.km / u.s) new_state = RVState(state.attractor, (r, v), state.plane) return new_state def propagate_many(self, state, tofs): state = state.to_vectors() - - rrs, vvs = cowell_vb( - state.attractor.k.to_value(u.km**3 / u.s**2), - *state.to_value(), - tofs.to_value(u.s), - self._rtol, - events=self._events, - func=self._func, + # TODO make sure tofs is sorted + + r0, v0 = state.to_value() + ( # pylint: disable=E0633,E1120 + _, + _, + _, + last_ts, + status, + t_idx, + rrs, + vvs, + ) = self._cowell_gf( # pylint: disable=E0633,E1120 + tofs.to_value(u.s), # tofs + r0, # rr + v0, # vv + state.attractor.k.to_value(u.km**3 / u.s**2), # argk + self._rtol, # rtol + self._atol, # atol + self._terminals, # event_terminals + self._directions, # event_directions ) + assert np.all((status == SOLVE_FINISHED) | (status == SOLVE_TERMINATED)) + + for last_t, event in zip(last_ts, self._events): + event.last_t_raw = last_t + # TODO: This should probably return a RVStateArray instead, # see discussion at https://github.com/poliastro/poliastro/pull/1492 return ( - rrs << u.km, - vvs << (u.km / u.s), + rrs[:t_idx, :] << u.km, + vvs[:t_idx, :] << (u.km / u.s), ) From 69405e1e3ee18a071ac4a6bcb1e939761dde25f9 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 5 Feb 2024 12:27:34 +0100 Subject: [PATCH 283/346] typing --- src/hapsira/core/propagation/cowell.py | 38 ++++++++++++-------------- 1 file changed, 18 insertions(+), 20 deletions(-) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index b18b26aeb..0184d087c 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -1,8 +1,6 @@ from math import isnan, nan from typing import Callable, Tuple -from numpy import ndarray - from ..jit import gjit, array_to_V_hf from ..math.ieee754 import EPS from ..math.ivp import ( @@ -81,25 +79,25 @@ def cowell_vb( "void(f[:],f[:],f[:],f,f,f,b1[:],f[:],f[:],f[:],f[:],f[:],i8[:],i8[:],f[:,:],f[:,:])", "(n),(m),(m),(),(),(),(o),(o)->(o),(o),(o),(o),(),(),(n,m),(n,m)", cache=False, - ) # n: tofs, m: dims, o: events + ) def cowell_gf( - tofs: ndarray, - rr: Tuple[float, float, float], - vv: Tuple[float, float, float], - argk: float, - rtol: float, - atol: float, - event_terminals: ndarray, - event_directions: ndarray, - event_g_olds: ndarray, # (out) - event_g_news: ndarray, # (out) - event_actives: ndarray, # out - event_last_ts: ndarray, # out - status: int, # out - t_idx: int, # out - rrs: ndarray, # out - vvs: ndarray, # out - ): # -> void(..., rrs, vvs, success) + tofs, + rr, + vv, + argk, + rtol, + atol, + event_terminals, + event_directions, + event_g_olds, + event_g_news, + event_actives, + event_last_ts, + status, + t_idx, + rrs, + vvs, + ): """ Solve an initial value problem for a system of ODEs. From 6e365a753c0017258eb4923a9a9f72a7b925f01f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 6 Feb 2024 10:56:14 +0100 Subject: [PATCH 284/346] fix name --- src/hapsira/core/propagation/cowell.py | 4 ++-- src/hapsira/twobody/propagation/cowell.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 0184d087c..303c36ed4 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -34,7 +34,7 @@ __all__ = [ - "cowell_vb", + "cowell_gb", "SOLVE_BRENTQFAILED", "SOLVE_FAILED", "SOLVE_RUNNING", @@ -50,7 +50,7 @@ SOLVE_TERMINATED = 1 -def cowell_vb( +def cowell_gb( events: Tuple = tuple(), func: Callable = func_twobody_hf, ) -> Callable: diff --git a/src/hapsira/twobody/propagation/cowell.py b/src/hapsira/twobody/propagation/cowell.py index 742d179fd..ab6302552 100644 --- a/src/hapsira/twobody/propagation/cowell.py +++ b/src/hapsira/twobody/propagation/cowell.py @@ -4,7 +4,7 @@ import numpy as np from hapsira.core.math.ieee754 import float_ -from hapsira.core.propagation.cowell import cowell_vb, SOLVE_FINISHED, SOLVE_TERMINATED +from hapsira.core.propagation.cowell import cowell_gb, SOLVE_FINISHED, SOLVE_TERMINATED from hapsira.core.propagation.base import func_twobody_hf from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import RVState @@ -35,7 +35,7 @@ def __init__(self, rtol=1e-11, atol=1e-12, events=tuple(), f=func_twobody_hf): self._events = events self._terminals = np.array([event.terminal for event in events], dtype=bool) self._directions = np.array([event.direction for event in events], dtype=float_) - self._cowell_gf = cowell_vb(events=events, func=f) + self._cowell_gf = cowell_gb(events=events, func=f) def propagate(self, state, tof): state = state.to_vectors() From 406934d20873ea2c60c6041099d30b871dd95eb2 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 6 Feb 2024 11:13:28 +0100 Subject: [PATCH 285/346] fix call into brentq_hf to gf --- src/hapsira/core/math/ivp/__init__.py | 4 ++-- src/hapsira/core/math/ivp/_brentq.py | 33 +++++++++++++++++++++++---- src/hapsira/threebody/restricted.py | 22 ++++++++++-------- 3 files changed, 43 insertions(+), 16 deletions(-) diff --git a/src/hapsira/core/math/ivp/__init__.py b/src/hapsira/core/math/ivp/__init__.py index 7d1562754..6af349f72 100644 --- a/src/hapsira/core/math/ivp/__init__.py +++ b/src/hapsira/core/math/ivp/__init__.py @@ -6,7 +6,7 @@ BRENTQ_XTOL, BRENTQ_RTOL, BRENTQ_MAXITER, - brentq_hf, + brentq_gb, brentq_dense_hf, ) from ._const import DENSE_SIG @@ -58,7 +58,7 @@ "DOP853_T_OLD", "DOP853_VV", "DOP853_VV_OLD", - "brentq_hf", + "brentq_gb", "brentq_dense_hf", "dispatcher_hb", "dop853_dense_interp_brentq_hb", diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index d68d3b197..f21d621c6 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -1,8 +1,9 @@ from math import fabs, isnan, nan +from typing import Callable from ._const import DENSE_SIG from ..ieee754 import EPS -from ...jit import hjit +from ...jit import hjit, gjit __all__ = [ @@ -13,7 +14,7 @@ "BRENTQ_XTOL", "BRENTQ_RTOL", "BRENTQ_MAXITER", - "brentq_hf", + "brentq_gb", "brentq_dense_hf", ] @@ -38,8 +39,8 @@ def _signbit_s_hf(a): return a < 0 -@hjit("Tuple([f,i8])(F(f(f)),f,f,f,f,f)", forceobj=True, nopython=False, cache=False) -def brentq_hf( +@hjit("Tuple([f,i8])(F(f(f)),f,f,f,f,f)") +def _brentq_hf( func, # callback_type xa, # double xb, # double @@ -132,6 +133,30 @@ def brentq_hf( return xcur, BRENTQ_CONVERR +def brentq_gb(func: Callable) -> Callable: + """ + Builds vectorized brentq + """ + + @gjit( + "void(f,f,f,f,f,f[:],i8[:])", + "(),(),(),(),()->(),()", + cache=False, + ) + def brentq_gf( + xa, + xb, + xtol, + rtol, + maxiter, + xcur, + status, + ): + xcur[0], status[0] = _brentq_hf(func, xa, xb, xtol, rtol, maxiter) + + return brentq_gf + + @hjit(f"Tuple([f,f,i8])(F(f(i8,f,{DENSE_SIG:s},f)),i8,f,f,f,f,f,{DENSE_SIG:s},f)") def brentq_dense_hf( func, # callback_type diff --git a/src/hapsira/threebody/restricted.py b/src/hapsira/threebody/restricted.py index 1addf91e2..fc7dda17b 100644 --- a/src/hapsira/threebody/restricted.py +++ b/src/hapsira/threebody/restricted.py @@ -9,7 +9,7 @@ from hapsira.core.jit import hjit from hapsira.core.math.ivp import ( - brentq_hf, + brentq_gb, BRENTQ_XTOL, BRENTQ_RTOL, BRENTQ_MAXITER, @@ -51,29 +51,31 @@ def eq_L123(xi): aux -= xi return aux + brentq_gf = brentq_gb(eq_L123) + lp = np.zeros((5,)) # L1 tol = 1e-11 # `brentq` uses a xtol of 2e-12, so it should be covered a = -pi2 + tol b = 1 - pi2 - tol - xi, status = brentq_hf( - eq_L123, a, b, BRENTQ_XTOL, BRENTQ_RTOL, BRENTQ_MAXITER - ) # TODO call into hf + xi, status = brentq_gf( # pylint: disable=E0633,E1120 + a, b, BRENTQ_XTOL, BRENTQ_RTOL, BRENTQ_MAXITER + ) assert status == BRENTQ_CONVERGED lp[0] = xi + pi2 # L2 - xi, status = brentq_hf( - eq_L123, 1, 1.5, BRENTQ_XTOL, BRENTQ_RTOL, BRENTQ_MAXITER - ) # TODO call into hf + xi, status = brentq_gf( # pylint: disable=E0633,E1120 + 1, 1.5, BRENTQ_XTOL, BRENTQ_RTOL, BRENTQ_MAXITER + ) assert status == BRENTQ_CONVERGED lp[1] = xi + pi2 # L3 - xi, status = brentq_hf( - eq_L123, -1.5, -1, BRENTQ_XTOL, BRENTQ_RTOL, BRENTQ_MAXITER - ) # TODO call into hf + xi, status = brentq_gf( # pylint: disable=E0633,E1120 + -1.5, -1, BRENTQ_XTOL, BRENTQ_RTOL, BRENTQ_MAXITER + ) assert status == BRENTQ_CONVERGED lp[2] = xi + pi2 From f60f96a46647cdbe4b033254973774a46c9ac470 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 6 Feb 2024 12:53:54 +0100 Subject: [PATCH 286/346] vectorized eclipse --- src/hapsira/core/events.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index 336a21f98..069e2cf3e 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -10,13 +10,18 @@ __all__ = [ + "ECLIPSE_UMBRA", "eclipse_function_hf", + "eclipse_function_gf", "line_of_sight_hf", "line_of_sight_gf", "elevation_function", ] +ECLIPSE_UMBRA = True + + @hjit("f(f,V,V,V,f,f,b1)") def eclipse_function_hf(k, rr, vv, r_sec, R_sec, R_primary, umbra): """Calculates a continuous shadow function. @@ -71,6 +76,26 @@ def eclipse_function_hf(k, rr, vv, r_sec, R_sec, R_primary, umbra): return shadow_function +@gjit( + "void(f,f[:],f[:],f[:],f,f,b1,f[:])", + "(),(n),(n),(n),(),(),()->()", +) +def eclipse_function_gf(k, rr, vv, r_sec, R_sec, R_primary, umbra, eclipse): + """ + Vectorized eclipse_function + """ + + eclipse[0] = eclipse_function_hf( + k, + array_to_V_hf(rr), + array_to_V_hf(vv), + array_to_V_hf(r_sec), + R_sec, + R_primary, + umbra, + ) + + @hjit("f(V,V,f)") def line_of_sight_hf(r1, r2, R): """Calculates the line of sight condition between two position vectors, r1 and r2. From 99ffacb4f66521b96cc4be8dfaa7d002443eb3ac Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 6 Feb 2024 12:54:21 +0100 Subject: [PATCH 287/346] run on new event implementation --- docs/source/examples/detecting-events.myst.md | 96 +++++++++---------- 1 file changed, 45 insertions(+), 51 deletions(-) diff --git a/docs/source/examples/detecting-events.myst.md b/docs/source/examples/detecting-events.myst.md index 6af490c50..ec41a376a 100644 --- a/docs/source/examples/detecting-events.myst.md +++ b/docs/source/examples/detecting-events.myst.md @@ -4,7 +4,7 @@ jupytext: extension: .md format_name: myst format_version: 0.13 - jupytext_version: 1.14.0 + jupytext_version: 1.16.0 kernelspec: display_name: Python 3 (ipykernel) language: python @@ -30,7 +30,7 @@ an event during an orbit's propagation is fairly simple: 2. Pass the `Event` object(s) as an argument to `CowellPropagator`. 3. Detect events! Optionally, the `terminal` and `direction` attributes can be set as required. -```{code-cell} +```{code-cell} ipython3 # Imports import numpy as np from numpy.linalg import norm @@ -63,11 +63,12 @@ from hapsira.util import time_range ## Altitude Crossing Event Let's define some natural perturbation conditions for our orbit so that its altitude decreases with time. -```{code-cell} +```{code-cell} ipython3 from hapsira.constants import H0_earth, rho0_earth -from hapsira.core.jit import array_to_V_hf +from hapsira.core.jit import djit +from hapsira.core.math.linalg import add_VV_hf from hapsira.core.perturbations import atmospheric_drag_exponential_hf -from hapsira.core.propagation import func_twobody +from hapsira.core.propagation.base import func_twobody_hf R = Earth.R.to_value(u.km) @@ -81,19 +82,18 @@ A_over_m = ((np.pi / 4.0) * (u.m**2) / (100 * u.kg)).to_value( rho0 = rho0_earth.to_value(u.kg / u.km**3) # kg/km^3 H0 = H0_earth.to_value(u.km) # km - -def f(t0, u_, k): - du_kep = func_twobody(t0, u_, k) - ax, ay, az = atmospheric_drag_exponential_hf( - t0, array_to_V_hf(u_[:3]), array_to_V_hf(u_[3:]), k, R=R, C_D=C_D, A_over_m=A_over_m, H0=H0, rho0=rho0 +@djit +def f_hf(t0, r0, v0, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, r0, v0, k) + a = atmospheric_drag_exponential_hf( + t0, r0, v0, k, R=R, C_D=C_D, A_over_m=A_over_m, H0=H0, rho0=rho0 ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_vv, a) ``` We shall use the `CowellPropagator` with the above perturbating conditions and pass the events we want to keep track of, in this case only the `AltitudeCrossEvent`. -```{code-cell} +```{code-cell} ipython3 tofs = np.arange(0, 2400, 100) << u.s orbit = Orbit.circular(Earth, 150 * u.km) @@ -102,7 +102,7 @@ thresh_alt = 50 # in km altitude_cross_event = AltitudeCrossEvent(thresh_alt, R) # Set up the event. events = [altitude_cross_event] -method = CowellPropagator(events=events, f=f) +method = CowellPropagator(events=events, f=f_hf) rr, _ = orbit.to_ephem( EpochsArray(orbit.epoch + tofs, method=method), ).rv() @@ -114,7 +114,7 @@ print( Let's see how did the orbit's altitude vary with time: -```{code-cell} +```{code-cell} ipython3 altitudes = np.apply_along_axis( norm, 1, (rr << u.km).value ) - Earth.R.to_value(u.km) @@ -131,7 +131,7 @@ Refer to the API documentation of the events to check the default values for `te Similar to the `AltitudeCrossEvent`, just pass the threshold latitude while instantiating the event. -```{code-cell} +```{code-cell} ipython3 orbit = Orbit.from_classical( Earth, 6900 << u.km, @@ -143,7 +143,7 @@ orbit = Orbit.from_classical( ) ``` -```{code-cell} +```{code-cell} ipython3 thresh_lat = 35 << u.deg latitude_cross_event = LatitudeCrossEvent(orbit, thresh_lat, terminal=True) events = [latitude_cross_event] @@ -158,16 +158,14 @@ print( Let's plot the latitude varying with time: -```{code-cell} -from hapsira.core.spheroid_location import cartesian_to_ellipsoidal +```{code-cell} ipython3 +from hapsira.core.spheroid_location import cartesian_to_ellipsoidal_gf -latitudes = [] -for r in rr: - position_on_body = (r / norm(r)) * Earth.R - _, lat, _ = cartesian_to_ellipsoidal( - Earth.R, Earth.R_polar, *position_on_body - ) - latitudes.append(np.rad2deg(lat)) +position_on_body = (rr.to_value(u.km) / norm(rr.to_value(u.km), axis = 1)[:, None]) * Earth.R.to_value(u.km) +_, latitudes, _ = cartesian_to_ellipsoidal_gf( + Earth.R.to_value(u.km), Earth.R_polar.to_value(u.km), *position_on_body.T +) +latitudes = np.rad2deg(latitudes) plt.plot(tofs[: len(rr)].to_value(u.s), latitudes) plt.title("Latitude variation") plt.ylabel("Latitude (in degrees)") @@ -179,7 +177,7 @@ The orbit's latitude would not change after the event was detected since we had Since the attractor is `Earth`, we could use `GroundtrackPlotter` for showing the groundtrack of the orbit on Earth. -```{code-cell} +```{code-cell} ipython3 from hapsira.earth import EarthSatellite from hapsira.earth.plotting import GroundtrackPlotter from hapsira.plotting import OrbitPlotter @@ -209,7 +207,7 @@ gp.plot( Viewing it in the `orthographic` projection mode, -```{code-cell} +```{code-cell} ipython3 gp.update_geos(projection_type="orthographic") gp.fig.show() ``` @@ -221,9 +219,7 @@ and voila! The groundtrack terminates almost at the 35 degree latitude mark. Users can detect umbra/penumbra crossings using the `UmbraEvent` and `PenumbraEvent` event classes, respectively. As seen from the above examples, the procedure doesn't change much. -```{code-cell} -from hapsira.core.events import eclipse_function - +```{code-cell} ipython3 attractor = Earth tof = 2 * u.d # Classical orbital elements @@ -240,11 +236,12 @@ orbit = Orbit.from_classical(attractor, *coe) Let's search for umbra crossings. -```{code-cell} -umbra_event = UmbraEvent(orbit, terminal=True) +```{code-cell} ipython3 +tofs = np.arange(0, 600, 30) << u.s + +umbra_event = UmbraEvent(orbit, tof, terminal=True) events = [umbra_event] -tofs = np.arange(0, 600, 30) << u.s method = CowellPropagator(events=events) rr, vv = orbit.to_ephem(EpochsArray(orbit.epoch + tofs, method=method)).rv() print( @@ -256,7 +253,9 @@ print( Let us plot the eclipse functions' variation with time. -```{code-cell} +```{code-cell} ipython3 +from hapsira.core.events import eclipse_function_gf, ECLIPSE_UMBRA + k = Earth.k.to_value(u.km**3 / u.s**2) R_sec = Sun.R.to_value(u.km) R_pri = Earth.R.to_value(u.km) @@ -269,12 +268,7 @@ r_sec = ((r_sec_ssb - r_pri_ssb).xyz << u.km).value rr = (rr << u.km).value vv = (vv << u.km / u.s).value -eclipses = [] # List to store values of eclipse_function. -for i in range(len(rr)): - r = rr[i] - v = vv[i] - eclipse = eclipse_function(k, np.hstack((r, v)), r_sec, R_sec, R_pri) - eclipses.append(eclipse) +eclipses = eclipse_function_gf(k, rr, vv, r_sec, R_sec, R_pri, ECLIPSE_UMBRA) plt.xlabel("Time (s)") plt.ylabel("Eclipse function") @@ -286,7 +280,7 @@ plt.plot(tofs[: len(rr)].to_value(u.s), eclipses) We could get some geometrical insights by plotting the orbit: -```{code-cell} +```{code-cell} ipython3 # Plot `Earth` at the instant of event occurence. Earth.plot( orbit.epoch.tdb + umbra_event.last_t, @@ -312,7 +306,7 @@ It seems our satellite is exiting the umbra region, as is evident from the orang This event detector aims to check for ascending and descending node crossings. Note that it could yield inaccurate results if the orbit is near-equatorial. -```{code-cell} +```{code-cell} ipython3 r = [-3182930.668, 94242.56, -85767.257] << u.km v = [505.848, 942.781, 7435.922] << u.km / u.s orbit = Orbit.from_vectors(Earth, r, v) @@ -320,13 +314,13 @@ orbit = Orbit.from_vectors(Earth, r, v) As a sanity check, let's check the orbit's inclination to ensure it is not near-zero: -```{code-cell} +```{code-cell} ipython3 print(orbit.inc) ``` Indeed, it isn't! -```{code-cell} +```{code-cell} ipython3 node_event = NodeCrossEvent(terminal=True) events = [node_event] @@ -339,7 +333,7 @@ print(f"The nodal cross time was {node_event.last_t} after the orbit's epoch") The plot below shows us the variation of the z coordinate of the orbit's position vector with time: -```{code-cell} +```{code-cell} ipython3 z_coords = [r[-1].to_value(u.km) for r in rr] plt.xlabel("Time (s)") plt.ylabel("Z coordinate of the position vector") @@ -349,7 +343,7 @@ plt.plot(tofs[: len(rr)].to_value(u.s), z_coords) We could do the same plotting done in `LatitudeCrossEvent` to check for equatorial crossings: -```{code-cell} +```{code-cell} ipython3 es = EarthSatellite(orbit, None) # Show the groundtrack plot from @@ -375,7 +369,7 @@ gp.plot( ) ``` -```{code-cell} +```{code-cell} ipython3 gp.update_geos(projection_type="orthographic") gp.fig.show() ``` @@ -389,7 +383,7 @@ either of the two crossings, the `direction` attribute is at our disposal! If we would like to track multiple events while propagating an orbit, we just need to add the concerned events inside `events`. Below, we show the case where `NodeCrossEvent` and `LatitudeCrossEvent` events are to be detected. -```{code-cell} +```{code-cell} ipython3 # NodeCrossEvent is detected earlier than the LatitudeCrossEvent. r = [-6142438.668, 3492467.56, -25767.257] << u.km v = [505.848, 942.781, 7435.922] << u.km / u.s @@ -407,9 +401,9 @@ tofs = [1, 2, 4, 6, 8, 10, 12] << u.s method = CowellPropagator(events=events) rr, vv = orbit.to_ephem(EpochsArray(orbit.epoch + tofs, method=method)).rv() -print(f"Node cross event termination time: {node_cross_event.last_t} s") +print(f"Node cross event termination time: {node_cross_event.last_t}") print( - f"Latitude cross event termination time: {latitude_cross_event.last_t} s" + f"Latitude cross event termination time: {latitude_cross_event.last_t}" ) ``` From 4fd153783e1c56765c0bff492b82c9e9f654e3ee Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 11 Feb 2024 15:14:48 +0100 Subject: [PATCH 288/346] force obj option --- src/hapsira/settings.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/hapsira/settings.py b/src/hapsira/settings.py index a18e06060..71ba3e024 100644 --- a/src/hapsira/settings.py +++ b/src/hapsira/settings.py @@ -133,6 +133,12 @@ def __init__(self): True, ) ) + self._add( + Setting( + "FORCEOBJ", + False, + ) + ), self._add( Setting( "PRECISION", From a122014c31974b1564bac31afdeced34da06e161 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 11 Feb 2024 15:16:44 +0100 Subject: [PATCH 289/346] force obj option --- src/hapsira/core/jit.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 913d788b3..6b7875303 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -115,9 +115,11 @@ def wrapper(inner_func: Callable) -> Callable: cache=settings["CACHE"].value, ) else: + assert settings["NOPYTHON"].value != settings["FORCEOBJ"].value wjit = nb.jit cfg = dict( nopython=settings["NOPYTHON"].value, + forceobj=settings["FORCEOBJ"].value, inline="always" if inline else "never", cache=settings["CACHE"].value, ) @@ -203,7 +205,9 @@ def wrapper(inner_func: Callable) -> Callable: cache=settings["CACHE"].value, ) if settings["TARGET"].value != "cuda": + assert settings["NOPYTHON"].value != settings["FORCEOBJ"].value cfg["nopython"] = settings["NOPYTHON"].value + cfg["forceobj"] = settings["FORCEOBJ"].value cfg.update(kwargs) logger.debug( @@ -249,7 +253,9 @@ def wrapper(inner_func: Callable) -> Callable: cache=settings["CACHE"].value, ) if settings["TARGET"].value != "cuda": + assert settings["NOPYTHON"].value != settings["FORCEOBJ"].value cfg["nopython"] = settings["NOPYTHON"].value + cfg["forceobj"] = settings["FORCEOBJ"].value cfg.update(kwargs) logger.debug( @@ -290,8 +296,10 @@ def wrapper(inner_func: Callable) -> Callable: Applies JIT """ + assert settings["NOPYTHON"].value != settings["FORCEOBJ"].value cfg = dict( nopython=settings["NOPYTHON"].value, + forceobj=settings["FORCEOBJ"].value, inline="always" if settings["INLINE"].value else "never", **kwargs, ) From d35c101a38537388c641f84a207edd0d21bb6fdc Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 13 Feb 2024 15:20:56 +0100 Subject: [PATCH 290/346] draft fix --- ...tural-and-artificial-perturbations.myst.md | 140 +++++++++--------- 1 file changed, 67 insertions(+), 73 deletions(-) diff --git a/docs/source/examples/natural-and-artificial-perturbations.myst.md b/docs/source/examples/natural-and-artificial-perturbations.myst.md index 4b14abeb9..dee44bcb3 100644 --- a/docs/source/examples/natural-and-artificial-perturbations.myst.md +++ b/docs/source/examples/natural-and-artificial-perturbations.myst.md @@ -4,7 +4,7 @@ jupytext: extension: .md format_name: myst format_version: 0.13 - jupytext_version: 1.14.1 + jupytext_version: 1.16.0 kernelspec: display_name: Python 3 (ipykernel) language: python @@ -25,13 +25,14 @@ from hapsira.bodies import Earth, Moon from hapsira.constants import rho0_earth, H0_earth from hapsira.core.elements import rv2coe_gf, RV2COE_TOL -from hapsira.core.jit import array_to_V_hf +from hapsira.core.jit import array_to_V_hf, djit, hjit +from hapsira.core.math.linalg import add_VV_hf from hapsira.core.perturbations import ( atmospheric_drag_exponential_hf, third_body_hf, J2_perturbation_hf, ) -from hapsira.core.propagation import func_twobody +from hapsira.core.propagation.base import func_twobody_hf from hapsira.ephem import build_ephem_interpolant from hapsira.plotting import OrbitPlotter from hapsira.plotting.orbit.backends import Plotly3D @@ -65,12 +66,13 @@ H0 = H0_earth.to(u.km).value tofs = TimeDelta(np.linspace(0 * u.h, 100000 * u.s, num=2000)) - -def f(t0, state, k): - du_kep = func_twobody(t0, state, k) - ax, ay, az = atmospheric_drag_exponential_hf( +@djit +def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + a = atmospheric_drag_exponential_hf( t0, - array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), + rr, + vv, k, R=R, C_D=C_D, @@ -78,13 +80,10 @@ def f(t0, state, k): H0=H0, rho0=rho0, ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - - return du_kep + du_ad - + return du_kep_rr, add_VV_hf(du_kep_vv, a) rr, _ = orbit.to_ephem( - EpochsArray(orbit.epoch + tofs, method=CowellPropagator(f=f)), + EpochsArray(orbit.epoch + tofs, method=CowellPropagator(f=f_hf)), ).rv() ``` @@ -117,7 +116,7 @@ events = [lithobrake_event] rr, _ = orbit.to_ephem( EpochsArray( - orbit.epoch + tofs, method=CowellPropagator(f=f, events=events) + orbit.epoch + tofs, method=CowellPropagator(f=f_hf, events=events) ), ).rv() @@ -145,20 +144,20 @@ v0 = np.array([-7.36138, -2.98997, 1.64354]) * u.km / u.s orbit = Orbit.from_vectors(Earth, r0, v0) tofs = TimeDelta(np.linspace(0, 48.0 * u.h, num=2000)) - - -def f(t0, state, k): - du_kep = func_twobody(t0, state, k) - ax, ay, az = J2_perturbation_hf( - t0, array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), k, J2=Earth.J2.value, R=Earth.R.to(u.km).value +_J2 = Earth.J2.value +_R = Earth.R.to(u.km).value + +@djit +def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + a = J2_perturbation_hf( + t0, rr, vv, k, J2=_J2, R=_R ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_vv, a) rr, vv = orbit.to_ephem( - EpochsArray(orbit.epoch + tofs, method=CowellPropagator(f=f)), + EpochsArray(orbit.epoch + tofs, method=CowellPropagator(f=f_hf)), ).rv() # This will be easier to compute when this is solved: @@ -206,26 +205,25 @@ initial = Orbit.from_classical( ) tofs = TimeDelta(np.linspace(0, 60 * u.day, num=1000)) +_moon_k = Moon.k.to(u.km**3 / u.s**2).value - -def f(t0, state, k): - du_kep = func_twobody(t0, state, k) - ax, ay, az = third_body_hf( +@djit(cache = False) +def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + a = third_body_hf( t0, - array_to_V_hf(state[:3]), - array_to_V_hf(state[3:]), + rr, + vv, k, - k_third=400 * Moon.k.to(u.km**3 / u.s**2).value, + k_third=400 * _moon_k, perturbation_body=body_r, ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_vv, a) # multiply Moon gravity by 400 so that effect is visible :) ephem = initial.to_ephem( - EpochsArray(initial.epoch + tofs, method=CowellPropagator(rtol=1e-6, f=f)), + EpochsArray(initial.epoch + tofs, method=CowellPropagator(rtol=1e-6, f=f_hf)), ) ``` @@ -264,24 +262,21 @@ orb0 = Orbit.from_classical( a_d_hf, _, t_f = change_ecc_inc(orb0, ecc_f, inc_f, f) - -def f(t0, state, k): - du_kep = func_twobody(t0, state, k) - ax, ay, az = a_d_hf( +@djit +def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + a = a_d_hf( t0, - array_to_V_hf(state[:3]), - array_to_V_hf(state[3:]), + rr, + vv, k, ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - - return du_kep + du_ad - + return du_kep_rr, add_VV_hf(du_kep_vv, a) tofs = TimeDelta(np.linspace(0, t_f, num=1000)) ephem2 = orb0.to_ephem( - EpochsArray(orb0.epoch + tofs, method=CowellPropagator(rtol=1e-6, f=f)), + EpochsArray(orb0.epoch + tofs, method=CowellPropagator(rtol=1e-6, f=f_hf)), ) ``` @@ -297,51 +292,52 @@ frame.plot_ephem(ephem2, label="orbit with artificial thrust") It might be of interest to determine what effect multiple perturbations have on a single object. In order to add multiple perturbations we can create a custom function that adds them up: ```{code-cell} ipython3 -from numba import njit as jit - -# Add @jit for speed! -@jit -def a_d(t0, state, k, J2, R, C_D, A_over_m, H0, rho0): - return np.array(J2_perturbation_hf(t0, array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), k, J2, R)) + np.array(atmospheric_drag_exponential_hf( - t0, array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), k, R, C_D, A_over_m, H0, rho0 - )) +@hjit("V(f,V,V,f,f,f,f,f,f,f)") +def a_d_hf(t0, rr, vv, k, J2, R, C_D, A_over_m, H0, rho0): + return add_VV_hf( + J2_perturbation_hf(t0, rr, vv, k, J2, R), + atmospheric_drag_exponential_hf( + t0, rr, vv, k, R, C_D, A_over_m, H0, rho0 + ) + ) ``` ```{code-cell} ipython3 # propagation times of flight and orbit tofs = TimeDelta(np.linspace(0, 10 * u.day, num=10 * 500)) orbit = Orbit.circular(Earth, 250 * u.km) # recall orbit from drag example +_J2 = Earth.J2.value - -def f(t0, state, k): - du_kep = func_twobody(t0, state, k) - ax, ay, az = a_d( +@djit +def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + a = a_d_hf( t0, - state, + rr, + vv, k, R=R, C_D=C_D, A_over_m=A_over_m, H0=H0, rho0=rho0, - J2=Earth.J2.value, + J2=_J2, ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - - return du_kep + du_ad - + return du_kep_rr, add_VV_hf(du_kep_rr, a) # propagate with J2 and atmospheric drag rr3, _ = orbit.to_ephem( - EpochsArray(orbit.epoch + tofs, method=CowellPropagator(f=f)), + EpochsArray(orbit.epoch + tofs, method=CowellPropagator(f=f_hf)), ).rv() -def f(t0, state, k): - du_kep = func_twobody(t0, state, k) - ax, ay, az = atmospheric_drag_exponential_hf( +@djit +def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + a = atmospheric_drag_exponential_hf( t0, - array_to_V_hf(state[:3]), array_to_V_hf(state[3:]), + rr, + vv, k, R=R, C_D=C_D, @@ -349,14 +345,12 @@ def f(t0, state, k): H0=H0, rho0=rho0, ) - du_ad = np.array([0, 0, 0, ax, ay, az]) - - return du_kep + du_ad + return du_kep_rr, add_VV_hf(du_kep_rr, a) # propagate with only atmospheric drag rr4, _ = orbit.to_ephem( - EpochsArray(orbit.epoch + tofs, method=CowellPropagator(f=f)), + EpochsArray(orbit.epoch + tofs, method=CowellPropagator(f=f_hf)), ).rv() ``` From 98eedfa17fc8d2a8e2a88ab60e3e9bc9fad4189d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 13 Feb 2024 15:23:49 +0100 Subject: [PATCH 291/346] rm nextafter; use fabs instead of abs --- src/hapsira/core/math/linalg.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 7855d82c6..dd2b1a531 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -1,6 +1,5 @@ -from math import inf, sqrt +from math import fabs, inf, sqrt -from .ieee754 import EPS from ..jit import hjit, vjit __all__ = [ @@ -16,7 +15,6 @@ "max_VV_hf", "mul_Vs_hf", "mul_VV_hf", - "nextafter_hf", "norm_V_hf", "norm_V_vf", "norm_VV_hf", @@ -28,7 +26,7 @@ @hjit("V(V)", inline=True) def abs_V_hf(x): - return abs(x[0]), abs(x[1]), abs(x[2]) + return fabs(x[0]), fabs(x[1]), fabs(x[2]) @hjit("V(V,f)", inline=True) @@ -124,13 +122,6 @@ def mul_VV_hf(a, b): return a[0] * b[0], a[1] * b[1], a[2] * b[2] -@hjit("f(f,f)") -def nextafter_hf(x, direction): - if x < direction: - return x + EPS - return x - EPS - - @hjit("f(V)", inline=True) def norm_V_hf(a): return sqrt(matmul_VV_hf(a, a)) From df4e1c1ab26afc73dc25d3924321fa916701b401 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 13 Feb 2024 15:24:22 +0100 Subject: [PATCH 292/346] temp fix: use nextafter from numpy, wait for math impl in numba --- src/hapsira/core/math/ieee754.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/hapsira/core/math/ieee754.py b/src/hapsira/core/math/ieee754.py index 677925f99..53acbb8b9 100644 --- a/src/hapsira/core/math/ieee754.py +++ b/src/hapsira/core/math/ieee754.py @@ -3,6 +3,7 @@ float32 as f4, float16 as f2, finfo, + nextafter, ) from ...settings import settings @@ -14,6 +15,7 @@ "f4", "f2", "float_", + "nextafter", ] From fab9b820cb84e88c4f2b55d6e9bb7619ca31dd62 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 13 Feb 2024 15:24:52 +0100 Subject: [PATCH 293/346] add note on nextafter --- src/hapsira/core/math/ieee754.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/math/ieee754.py b/src/hapsira/core/math/ieee754.py index 53acbb8b9..ae9cb3ec0 100644 --- a/src/hapsira/core/math/ieee754.py +++ b/src/hapsira/core/math/ieee754.py @@ -3,7 +3,7 @@ float32 as f4, float16 as f2, finfo, - nextafter, + nextafter, # TODO switch to math module ) from ...settings import settings From 600cee953054cae6f891c00cf49caf4048d2ed30 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Tue, 13 Feb 2024 15:27:56 +0100 Subject: [PATCH 294/346] use better nextafter; use fabs instead of abs --- src/hapsira/core/math/ivp/_rkstepimpl.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkstepimpl.py b/src/hapsira/core/math/ivp/_rkstepimpl.py index f39cb86cc..0e4935a29 100644 --- a/src/hapsira/core/math/ivp/_rkstepimpl.py +++ b/src/hapsira/core/math/ivp/_rkstepimpl.py @@ -1,9 +1,10 @@ -from math import inf +from math import inf, fabs from ._const import ERROR_EXPONENT, KSIG, MAX_FACTOR, MIN_FACTOR, SAFETY from ._rkstep import rk_step_hf from ._rkerror import estimate_error_norm_V_hf -from ..linalg import abs_V_hf, add_Vs_hf, max_VV_hf, mul_Vs_hf, nextafter_hf +from ..ieee754 import nextafter +from ..linalg import abs_V_hf, add_Vs_hf, max_VV_hf, mul_Vs_hf from ...jit import hjit, DSIG @@ -19,7 +20,7 @@ def step_impl_hf( fun, argk, t, rr, vv, fr, fv, rtol, atol, direction, h_abs, t_bound, K ): - min_step = 10 * abs(nextafter_hf(t, direction * inf) - t) + min_step = 10 * fabs(nextafter(t, direction * inf) - t) if h_abs < min_step: h_abs = min_step @@ -48,7 +49,7 @@ def step_impl_hf( t_new = t_bound h = t_new - t - h_abs = abs(h) + h_abs = fabs(h) rr_new, vv_new, fr_new, fv_new, K_new = rk_step_hf( fun, From 3d5c0d8c32e87dfe90d965dc218af8dee74d4600 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 10:25:49 +0100 Subject: [PATCH 295/346] rm pow; add missing export; add safe divs --- src/hapsira/core/math/linalg.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index dd2b1a531..488117693 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -9,6 +9,7 @@ "cross_VV_hf", "div_Vs_hf", "div_VV_hf", + "div_ss_hf", "matmul_MM_hf", "matmul_VM_hf", "matmul_VV_hf", @@ -60,12 +61,12 @@ def div_ss_hf(a, b): @hjit("V(V,V)", inline=True) def div_VV_hf(x, y): - return x[0] / y[0], x[1] / y[1], x[2] / y[2] + return div_ss_hf(x[0], y[0]), div_ss_hf(x[1], y[1]), div_ss_hf(x[2], y[2]) @hjit("V(V,f)", inline=True) def div_Vs_hf(v, s): - return v[0] / s, v[1] / s, v[2] / s + return div_ss_hf(v[0], s), div_ss_hf(v[1], s), div_ss_hf(v[2], s) @hjit("M(M,M)", inline=True) @@ -135,7 +136,7 @@ def norm_V_vf(a, b, c): @hjit("f(V,V)", inline=True) def norm_VV_hf(x, y): - return sqrt(x[0] ** 2 + x[1] ** 2 + x[2] ** 2 + y[0] ** 2 + y[1] ** 2 + y[2] ** 2) + return sqrt(matmul_VV_hf(x, x) + matmul_VV_hf(y, y)) @hjit("f(f)", inline=True) From fc1a73de4b3c395a45403d7b371e108a007b6d22 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 10:28:23 +0100 Subject: [PATCH 296/346] add safe divs; abs to fabs --- src/hapsira/core/math/ivp/_rkerror.py | 101 +++++++++++++------------- 1 file changed, 51 insertions(+), 50 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkerror.py b/src/hapsira/core/math/ivp/_rkerror.py index 5f3900156..98e04f6eb 100644 --- a/src/hapsira/core/math/ivp/_rkerror.py +++ b/src/hapsira/core/math/ivp/_rkerror.py @@ -1,8 +1,9 @@ -from math import sqrt +from math import fabs, sqrt from ._const import N_RV, KSIG from ._dop853_coefficients import E3 as _E3, E5 as _E5 from ..ieee754 import float_ +from ..linalg import div_ss_hf from ...jit import hjit @@ -20,7 +21,7 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): K00, K01, K02, K03, K04, K05, K06, K07, K08, K09, K10, K11, K12 = K err3 = ( - ( + div_ss_hf( K00[0] * E3[0] + K01[0] * E3[1] + K02[0] * E3[2] @@ -33,10 +34,10 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[0] * E3[9] + K10[0] * E3[10] + K11[0] * E3[11] - + K12[0] * E3[12] - ) - / scale_r[0], - ( + + K12[0] * E3[12], + scale_r[0], + ), + div_ss_hf( K00[1] * E3[0] + K01[1] * E3[1] + K02[1] * E3[2] @@ -49,10 +50,10 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[1] * E3[9] + K10[1] * E3[10] + K11[1] * E3[11] - + K12[1] * E3[12] - ) - / scale_r[1], - ( + + K12[1] * E3[12], + scale_r[1], + ), + div_ss_hf( K00[2] * E3[0] + K01[2] * E3[1] + K02[2] * E3[2] @@ -65,10 +66,10 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[2] * E3[9] + K10[2] * E3[10] + K11[2] * E3[11] - + K12[2] * E3[12] - ) - / scale_r[2], - ( + + K12[2] * E3[12], + scale_r[2], + ), + div_ss_hf( K00[3] * E3[0] + K01[3] * E3[1] + K02[3] * E3[2] @@ -81,10 +82,10 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[3] * E3[9] + K10[3] * E3[10] + K11[3] * E3[11] - + K12[3] * E3[12] - ) - / scale_v[0], - ( + + K12[3] * E3[12], + scale_v[0], + ), + div_ss_hf( K00[4] * E3[0] + K01[4] * E3[1] + K02[4] * E3[2] @@ -97,10 +98,10 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[4] * E3[9] + K10[4] * E3[10] + K11[4] * E3[11] - + K12[4] * E3[12] - ) - / scale_v[1], - ( + + K12[4] * E3[12], + scale_v[1], + ), + div_ss_hf( K00[5] * E3[0] + K01[5] * E3[1] + K02[5] * E3[2] @@ -113,12 +114,12 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[5] * E3[9] + K10[5] * E3[10] + K11[5] * E3[11] - + K12[5] * E3[12] - ) - / scale_v[2], + + K12[5] * E3[12], + scale_v[2], + ), ) err5 = ( - ( + div_ss_hf( K00[0] * E5[0] + K01[0] * E5[1] + K02[0] * E5[2] @@ -131,10 +132,10 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[0] * E5[9] + K10[0] * E5[10] + K11[0] * E5[11] - + K12[0] * E5[12] - ) - / scale_r[0], - ( + + K12[0] * E5[12], + scale_r[0], + ), + div_ss_hf( K00[1] * E5[0] + K01[1] * E5[1] + K02[1] * E5[2] @@ -147,10 +148,10 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[1] * E5[9] + K10[1] * E5[10] + K11[1] * E5[11] - + K12[1] * E5[12] - ) - / scale_r[1], - ( + + K12[1] * E5[12], + scale_r[1], + ), + div_ss_hf( K00[2] * E5[0] + K01[2] * E5[1] + K02[2] * E5[2] @@ -163,10 +164,10 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[2] * E5[9] + K10[2] * E5[10] + K11[2] * E5[11] - + K12[2] * E5[12] - ) - / scale_r[2], - ( + + K12[2] * E5[12], + scale_r[2], + ), + div_ss_hf( K00[3] * E5[0] + K01[3] * E5[1] + K02[3] * E5[2] @@ -179,10 +180,10 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[3] * E5[9] + K10[3] * E5[10] + K11[3] * E5[11] - + K12[3] * E5[12] - ) - / scale_v[0], - ( + + K12[3] * E5[12], + scale_v[0], + ), + div_ss_hf( K00[4] * E5[0] + K01[4] * E5[1] + K02[4] * E5[2] @@ -195,10 +196,10 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[4] * E5[9] + K10[4] * E5[10] + K11[4] * E5[11] - + K12[4] * E5[12] - ) - / scale_v[1], - ( + + K12[4] * E5[12], + scale_v[1], + ), + div_ss_hf( K00[5] * E5[0] + K01[5] * E5[1] + K02[5] * E5[2] @@ -211,9 +212,9 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): + K09[5] * E5[9] + K10[5] * E5[10] + K11[5] * E5[11] - + K12[5] * E5[12] - ) - / scale_v[2], + + K12[5] * E5[12], + scale_v[2], + ), ) err5_norm_2 = ( @@ -237,4 +238,4 @@ def estimate_error_norm_V_hf(K, h, scale_r, scale_v): return 0.0 denom = err5_norm_2 + 0.01 * err3_norm_2 - return abs(h) * err5_norm_2 / sqrt(denom * N_RV) + return fabs(h) * div_ss_hf(err5_norm_2, sqrt(denom * N_RV)) From 363932f37434c464508d1260112fe307563ec62b Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 10:32:07 +0100 Subject: [PATCH 297/346] clean up pows --- src/hapsira/core/propagation/base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/hapsira/core/propagation/base.py b/src/hapsira/core/propagation/base.py index fb3ee20b8..f96066ae2 100644 --- a/src/hapsira/core/propagation/base.py +++ b/src/hapsira/core/propagation/base.py @@ -1,3 +1,5 @@ +from math import pow as pow_ + from ..jit import djit @@ -24,6 +26,6 @@ def func_twobody_hf(t0, rr, vv, k): """ x, y, z = rr vx, vy, vz = vv - r3 = (x**2 + y**2 + z**2) ** 1.5 + r3 = pow_(x * x + y * y + z * z, 1.5) return (vx, vy, vz), (-k * x / r3, -k * y / r3, -k * z / r3) From 19eef134d7b87cc357f84f192a54f4a3773d5e4f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 10:35:26 +0100 Subject: [PATCH 298/346] rm todo --- src/hapsira/core/propagation/cowell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 303c36ed4..9a4c8ebdd 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -60,7 +60,7 @@ def cowell_gb( assert hasattr(func, "djit") # DEBUG check for compiler flag - EVENTS = len(events) # TODO compile as const + EVENTS = len(events) event_impl_hf = dispatcher_hb( funcs=tuple(event.impl_hf for event in events), From 83e691ae453f171551e2eb311d85e82fbbdf9b72 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 10:36:35 +0100 Subject: [PATCH 299/346] asserts --- src/hapsira/twobody/propagation/cowell.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/hapsira/twobody/propagation/cowell.py b/src/hapsira/twobody/propagation/cowell.py index ab6302552..62b44fd2b 100644 --- a/src/hapsira/twobody/propagation/cowell.py +++ b/src/hapsira/twobody/propagation/cowell.py @@ -4,7 +4,13 @@ import numpy as np from hapsira.core.math.ieee754 import float_ -from hapsira.core.propagation.cowell import cowell_gb, SOLVE_FINISHED, SOLVE_TERMINATED +from hapsira.core.propagation.cowell import ( + cowell_gb, + SOLVE_FINISHED, + SOLVE_TERMINATED, + SOLVE_BRENTQFAILED, + SOLVE_FAILED, +) from hapsira.core.propagation.base import func_twobody_hf from hapsira.twobody.propagation.enums import PropagatorKind from hapsira.twobody.states import RVState @@ -63,6 +69,8 @@ def propagate(self, state, tof): self._directions, # event_directions ) + assert np.all((status != SOLVE_FAILED)) + assert np.all((status != SOLVE_BRENTQFAILED)) assert np.all((status == SOLVE_FINISHED) | (status == SOLVE_TERMINATED)) for last_t, event in zip(last_ts, self._events): @@ -99,6 +107,8 @@ def propagate_many(self, state, tofs): self._directions, # event_directions ) + assert np.all((status != SOLVE_FAILED)) + assert np.all((status != SOLVE_BRENTQFAILED)) assert np.all((status == SOLVE_FINISHED) | (status == SOLVE_TERMINATED)) for last_t, event in zip(last_ts, self._events): From 8ca6d248c58c168a91701ad3385e2e7744be1676 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 11:01:54 +0100 Subject: [PATCH 300/346] safe pows --- src/hapsira/core/perturbations.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/perturbations.py b/src/hapsira/core/perturbations.py index 15e42aac8..4d154fe91 100644 --- a/src/hapsira/core/perturbations.py +++ b/src/hapsira/core/perturbations.py @@ -1,4 +1,4 @@ -from math import exp +from math import exp, pow as pow_ from .events import line_of_sight_hf from .jit import hjit @@ -48,9 +48,8 @@ def J2_perturbation_hf(t0, rr, vv, k, J2, R): """ r = norm_V_hf(rr) - factor = (3.0 / 2.0) * k * J2 * (R**2) / (r**5) - - a_base = 5.0 * rr[2] ** 2 / r**2 + factor = 1.5 * k * J2 * R * R / pow_(r, 5) + a_base = 5.0 * rr[2] * rr[2] / (r * r) a = a_base - 1, a_base - 1, a_base - 3 return mul_Vs_hf(mul_VV_hf(a, rr), factor) From 2808c9ae32181cf91fcf6408a89e0c336c4fdf9e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 11:15:43 +0100 Subject: [PATCH 301/346] fix pertubation calls --- ...tural-and-artificial-perturbations.myst.md | 27 +++++++++---------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/docs/source/examples/natural-and-artificial-perturbations.myst.md b/docs/source/examples/natural-and-artificial-perturbations.myst.md index dee44bcb3..1d1f5bb3d 100644 --- a/docs/source/examples/natural-and-artificial-perturbations.myst.md +++ b/docs/source/examples/natural-and-artificial-perturbations.myst.md @@ -316,21 +316,20 @@ def f_hf(t0, rr, vv, k): rr, vv, k, - R=R, - C_D=C_D, - A_over_m=A_over_m, - H0=H0, - rho0=rho0, - J2=_J2, + _J2, + R, + C_D, + A_over_m, + H0, + rho0, ) - return du_kep_rr, add_VV_hf(du_kep_rr, a) + return du_kep_rr, add_VV_hf(du_kep_vv, a) # propagate with J2 and atmospheric drag rr3, _ = orbit.to_ephem( EpochsArray(orbit.epoch + tofs, method=CowellPropagator(f=f_hf)), ).rv() - @djit def f_hf(t0, rr, vv, k): du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) @@ -339,13 +338,13 @@ def f_hf(t0, rr, vv, k): rr, vv, k, - R=R, - C_D=C_D, - A_over_m=A_over_m, - H0=H0, - rho0=rho0, + R, + C_D, + A_over_m, + H0, + rho0, ) - return du_kep_rr, add_VV_hf(du_kep_rr, a) + return du_kep_rr, add_VV_hf(du_kep_vv, a) # propagate with only atmospheric drag From 529ee3c991dfb5d247a38674a56596b3ed965a27 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 11:47:34 +0100 Subject: [PATCH 302/346] update notebook --- ...pagation-using-cowells-formulation.myst.md | 51 ++++++++++--------- 1 file changed, 26 insertions(+), 25 deletions(-) diff --git a/docs/source/examples/propagation-using-cowells-formulation.myst.md b/docs/source/examples/propagation-using-cowells-formulation.myst.md index 74c87eb21..b4a400d37 100644 --- a/docs/source/examples/propagation-using-cowells-formulation.myst.md +++ b/docs/source/examples/propagation-using-cowells-formulation.myst.md @@ -4,7 +4,7 @@ jupytext: extension: .md format_name: myst format_version: 0.13 - jupytext_version: 1.14.1 + jupytext_version: 1.16.0 kernelspec: display_name: Python 3 (ipykernel) language: python @@ -25,7 +25,7 @@ $$\ddot{\mathbb{r}} = -\frac{\mu}{|\mathbb{r}|^3} \mathbb{r} + \mathbb{a}_d$$ +++ -
An earlier version of this notebook allowed for more flexibility and interactivity, but was considerably more complex. Future versions of hapsira and plotly might bring back part of that functionality, depending on user feedback. You can still download the older version [here](https://github.com/hapsira/hapsira/blob/0.8.x/docs/source/examples/Propagation%20using%20Cowell's%20formulation.ipynb).
+
An earlier version of this notebook allowed for more flexibility and interactivity, but was considerably more complex. Future versions of hapsira and plotly might bring back part of that functionality, depending on user feedback. You can still download the older version [here](https://github.com/pleiszenburg/hapsira/blob/0.8.x/docs/source/examples/Propagation%20using%20Cowell's%20formulation.ipynb).
+++ @@ -40,7 +40,9 @@ from astropy import units as u import numpy as np from hapsira.bodies import Earth -from hapsira.core.propagation import func_twobody +from hapsira.core.jit import djit, hjit +from hapsira.core.math.linalg import add_VV_hf, mul_Vs_hf, norm_V_hf +from hapsira.core.propagation.base import func_twobody_hf from hapsira.examples import iss from hapsira.plotting import OrbitPlotter from hapsira.plotting.orbit.backends import Plotly3D @@ -57,22 +59,24 @@ accel = 2e-5 ``` ```{code-cell} ipython3 -def constant_accel_factory(accel): - def constant_accel(t0, u, k): - v = u[3:] - norm_v = (v[0] ** 2 + v[1] ** 2 + v[2] ** 2) ** 0.5 - return accel * v / norm_v +def constant_accel_hb(accel): - return constant_accel + @hjit("V(f,V,V,f)", cache = False) + def constant_accel_hf(t0, rr, vv, k): + norm_v = norm_V_hf(vv) + return mul_Vs_hf(vv, accel / norm_v) + + return constant_accel_hf ``` ```{code-cell} ipython3 -def f(t0, state, k): - du_kep = func_twobody(t0, state, k) - ax, ay, az = constant_accel_factory(accel)(t0, state, k) - du_ad = np.array([0, 0, 0, ax, ay, az]) +constant_accel_hf = constant_accel_hb(accel) - return du_kep + du_ad +@djit +def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + a = constant_accel_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, a) ``` ```{code-cell} ipython3 @@ -82,7 +86,7 @@ times ```{code-cell} ipython3 ephem = iss.to_ephem( - EpochsArray(iss.epoch + times, method=CowellPropagator(rtol=1e-11, f=f)), + EpochsArray(iss.epoch + times, method=CowellPropagator(rtol=1e-11, f=f_hf)), ) ``` @@ -165,18 +169,15 @@ So let's create a new circular orbit and perform the necessary checks, assuming orb = Orbit.circular(Earth, 500 << u.km) tof = 20 * orb.period -ad = constant_accel_factory(1e-7) - - -def f(t0, state, k): - du_kep = func_twobody(t0, state, k) - ax, ay, az = ad(t0, state, k) - du_ad = np.array([0, 0, 0, ax, ay, az]) - - return du_kep + du_ad +ad_hf = constant_accel_hb(1e-7) +@djit +def f_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + a = ad_hf(t0, rr, vv, k) + return du_kep_rr, add_VV_hf(du_kep_vv, a) -orb_final = orb.propagate(tof, method=CowellPropagator(f=f)) +orb_final = orb.propagate(tof, method=CowellPropagator(f=f_hf)) ``` ```{code-cell} ipython3 From 43014c6a9e975c007d61566aeffcd67e3d229379 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 16:02:06 +0100 Subject: [PATCH 303/346] fix arg; fix missing import --- docs/source/quickstart.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/source/quickstart.md b/docs/source/quickstart.md index 10033448e..1995fe8d1 100644 --- a/docs/source/quickstart.md +++ b/docs/source/quickstart.md @@ -155,6 +155,7 @@ To explore different propagation algorithms, check out the {py:mod}`hapsira.twob The `propagate` method gives you the final orbit at the epoch you designated. To retrieve the whole trajectory instead, you can use {py:meth}`hapsira.twobody.orbit.scalar.Orbit.to_ephem`, which returns an {{ Ephem }} instance: ```python +from astropy.time import Time from hapsira.twobody.sampling import EpochsArray, TrueAnomalyBounds, EpochBounds from hapsira.util import time_range @@ -165,7 +166,7 @@ end_date = Time("2022-07-11 07:05", scale="utc") ephem1 = iss.to_ephem() # Explicit times given -ephem2 = iss.to_ephem(strategy=EpochsArray(epochs=time_range(start_date, end_date))) +ephem2 = iss.to_ephem(strategy=EpochsArray(epochs=time_range(start_date, end=end_date))) # Automatic grid, true anomaly limits ephem3 = iss.to_ephem(strategy=TrueAnomalyBounds(min_nu=0 << u.deg, max_nu=180 << u.deg)) From 7dabc12e6ac541d9beff432acf1558ccca9df76f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 16:09:44 +0100 Subject: [PATCH 304/346] fix quick cowell --- docs/source/quickstart.md | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/docs/source/quickstart.md b/docs/source/quickstart.md index 1995fe8d1..a72040247 100644 --- a/docs/source/quickstart.md +++ b/docs/source/quickstart.md @@ -194,27 +194,27 @@ ephem4 = iss.to_ephem(strategy=EpochBounds(min_epoch=start_date, max_epoch=end_d Apart from the Keplerian propagators, hapsira also allows you to define custom perturbation accelerations to study non Keplerian orbits, thanks to Cowell's method: ```python ->>> from numba import njit >>> import numpy as np ->>> from hapsira.core.propagation import func_twobody +>>> from hapsira.core.jit import hjit, djit +>>> from hapsira.core.math.linalg import add_VV_hf, mul_Vs_hf, norm_V_hf +>>> from hapsira.core.propagation.base import func_twobody_hf >>> from hapsira.twobody.propagation import CowellPropagator >>> r0 = [-2384.46, 5729.01, 3050.46] << u.km >>> v0 = [-7.36138, -2.98997, 1.64354] << (u.km / u.s) >>> initial = Orbit.from_vectors(Earth, r0, v0) ->>> @njit -... def accel(t0, state, k): +>>> @hjit("V(f,V,V,f)") +... def accel_hf(t0, rr, vv, k): ... """Constant acceleration aligned with the velocity. """ -... v_vec = state[3:] -... norm_v = (v_vec * v_vec).sum() ** 0.5 -... return 1e-5 * v_vec / norm_v +... norm_v = norm_V_hf(vv) +... return mul_Vs_hf(vv, 1e-5 / norm_v) ... -... def f(t0, u_, k): -... du_kep = func_twobody(t0, u_, k) -... ax, ay, az = accel(t0, u_, k) -... du_ad = np.array([0, 0, 0, ax, ay, az]) -... return du_kep + du_ad +... @djit +... def f_hf(t0, rr, vv, k): +... du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) +... a = accel_hf(t0, rr, vv, k) +... return du_kep_rr, add_VV_hf(du_kep_vv, a) ->>> initial.propagate(3 << u.day, method=CowellPropagator(f=f)) +>>> initial.propagate(3 << u.day, method=CowellPropagator(f=f_hf)) 18255 x 21848 km x 28.0 deg (GCRS) orbit around Earth (♁) at epoch J2000.008 (TT) ``` From c0c694077623f038ced75dd01496f33e9367f11e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 16:14:28 +0100 Subject: [PATCH 305/346] fix quick pert --- docs/source/quickstart.md | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/docs/source/quickstart.md b/docs/source/quickstart.md index a72040247..e659a1057 100644 --- a/docs/source/quickstart.md +++ b/docs/source/quickstart.md @@ -221,17 +221,18 @@ Apart from the Keplerian propagators, hapsira also allows you to define custom p Some natural perturbations are available in hapsira to be used directly in this way. For instance, to examine the effect of J2 perturbation: ```python ->>> from hapsira.core.perturbations import J2_perturbation ->>> tofs = [48.0] << u.h ->>> def f(t0, u_, k): -... du_kep = func_twobody(t0, u_, k) -... ax, ay, az = J2_perturbation( -... t0, u_, k, J2=Earth.J2.value, R=Earth.R.to(u.km).value +>>> from hapsira.core.perturbations import J2_perturbation_hf +>>> tofs = 48.0 << u.h +>>> _J2, _R = Earth.J2.value, Earth.R.to(u.km).value +>>> @djit +... def f_hf(t0, rr, vv, k): +... du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) +... a = J2_perturbation_hf( +... t0, rr, vv, k, _J2, _R ... ) -... du_ad = np.array([0, 0, 0, ax, ay, az]) -... return du_kep + du_ad +... return du_kep_rr, add_VV_hf(du_kep_vv, a) ->>> final = initial.propagate(tofs, method=CowellPropagator(f=f)) +>>> final = initial.propagate(tofs, method=CowellPropagator(f=f_hf)) ``` The J2 perturbation changes the orbit parameters (from Curtis example 12.2): From dd9569f984f311c087a05922a27ffe22fac53de7 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 21:15:45 +0100 Subject: [PATCH 306/346] handle u.one --- src/hapsira/twobody/thrust/change_ecc_inc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/twobody/thrust/change_ecc_inc.py b/src/hapsira/twobody/thrust/change_ecc_inc.py index 9c437beab..4eb752c61 100644 --- a/src/hapsira/twobody/thrust/change_ecc_inc.py +++ b/src/hapsira/twobody/thrust/change_ecc_inc.py @@ -30,7 +30,7 @@ def change_ecc_inc(orb_0, ecc_f, inc_f, f): k=orb_0.attractor.k.to_value(u.km**3 / u.s**2), a=orb_0.a.to_value(u.km), ecc_0=orb_0.ecc.value, - ecc_f=ecc_f, + ecc_f=getattr(ecc_f, "value", ecc_f), # in case of u.one inc_0=orb_0.inc.to_value(u.rad), inc_f=inc_f.to_value(u.rad), argp=orb_0.argp.to_value(u.rad), From 2b87fc15ea389a2bd395b255f124ef7f3e36297f Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 21:20:10 +0100 Subject: [PATCH 307/346] f_geo --- docs/source/quickstart.md | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/docs/source/quickstart.md b/docs/source/quickstart.md index e659a1057..aabcb5c1e 100644 --- a/docs/source/quickstart.md +++ b/docs/source/quickstart.md @@ -249,6 +249,7 @@ The J2 perturbation changes the orbit parameters (from Curtis example 12.2): In addition to natural perturbations, hapsira also has built-in artificial perturbations (thrust guidance laws) aimed at intentional change of some orbital elements. For example, to simultaneously change eccentricity and inclination: ```python +>>> from hapsira.twobody.thrust import change_ecc_inc >>> ecc_0, ecc_f = [0.4, 0.0] << u.one >>> a = 42164 << u.km >>> inc_0 = 0.0 << u.deg # baseline @@ -262,20 +263,20 @@ In addition to natural perturbations, hapsira also has built-in artificial pertu ... a, ... ecc_0, ... inc_0, -... 0, +... 0 << u.deg, ... argp, -... 0, +... 0 << u.deg, ... ) ->>> a_d, _, t_f = change_ecc_inc(orb0, ecc_f, inc_f, f) +>>> a_d_hf, _, t_f = change_ecc_inc(orb0, ecc_f, inc_f, f) # Propagate orbit ->>> def f_geo(t0, u_, k): -... du_kep = func_twobody(t0, u_, k) -... ax, ay, az = a_d(t0, u_, k) -... du_ad = np.array([0, 0, 0, ax, ay, az]) -... return du_kep + du_ad +>>> @djit +... def f_geo_hf(t0, rr, vv, k): +... du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) +... a = a_d_hf(t0, rr, vv, k) +... return du_kep_rr, add_VV_hf(du_kep_vv, a) ->>> orbf = orb0.propagate(t_f << u.s, method=CowellPropagator(f=f_geo, rtol=1e-8)) +>>> orbf = orb0.propagate(t_f << u.s, method=CowellPropagator(f=f_geo_hf, rtol=1e-8)) ``` The thrust changes orbit parameters as desired (within errors): From 80dba2a47c4a9c0573cbe954edfaccdc64d78ca8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 21:33:07 +0100 Subject: [PATCH 308/346] fix plot; fix lambert --- docs/source/quickstart.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/source/quickstart.md b/docs/source/quickstart.md index aabcb5c1e..e89cfd278 100644 --- a/docs/source/quickstart.md +++ b/docs/source/quickstart.md @@ -349,8 +349,9 @@ To easily visualize several orbits in two dimensions, you can run this code: ```python from hapsira.plotting import OrbitPlotter +from hapsira.plotting.orbit.backends import Matplotlib2D -op = OrbitPlotter(backend_name="matplotlib2D") +op = OrbitPlotter(backend=Matplotlib2D()) orb_a, orb_f = orb_i.apply_maneuver(hoh, intermediate=True) op.plot(orb_i, label="Initial orbit") op.plot(orb_a, label="Transfer orbit") @@ -380,7 +381,7 @@ The {py:class}`hapsira.ephem.Ephem` class allows you to retrieve a planetary orb ```python >>> from astropy.time import Time ->>> epoch = time.Time("2020-04-29 10:43") # UTC by default +>>> epoch = Time("2020-04-29 10:43") # UTC by default >>> from hapsira.ephem import Ephem >>> earth = Ephem.from_body(Earth, epoch.tdb) >>> earth @@ -443,9 +444,9 @@ And these are the results: ```python >>> dv_a -(, ) +(, ) >>> dv_b -(, ) +(, ) ``` ```{figure} _static/msl.png From 49ad96d6d92dd0c0441d9f5fdeb74d5b636c2396 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 14 Feb 2024 21:37:34 +0100 Subject: [PATCH 309/346] trigger plot --- .../examples/propagation-using-cowells-formulation.myst.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/examples/propagation-using-cowells-formulation.myst.md b/docs/source/examples/propagation-using-cowells-formulation.myst.md index b4a400d37..3019474d9 100644 --- a/docs/source/examples/propagation-using-cowells-formulation.myst.md +++ b/docs/source/examples/propagation-using-cowells-formulation.myst.md @@ -97,6 +97,7 @@ frame = OrbitPlotter(backend=Plotly3D()) frame.set_attractor(Earth) frame.plot_ephem(ephem, label="ISS") +frame.show() ``` ## Error checking From 61b74c21628bdd93266713eccb7213999fb94529 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 19 Feb 2024 19:06:18 +0100 Subject: [PATCH 310/346] missing matmul op --- src/hapsira/core/math/linalg.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 488117693..7b0a419fa 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -11,6 +11,7 @@ "div_VV_hf", "div_ss_hf", "matmul_MM_hf", + "matmul_MV_hf", "matmul_VM_hf", "matmul_VV_hf", "max_VV_hf", @@ -99,6 +100,15 @@ def matmul_VM_hf(a, b): ) +@hjit("V(M,V)", inline=True) +def matmul_MV_hf(a, b): + return ( + b[0] * a[0][0] + b[1] * a[0][1] + b[2] * a[0][2], + b[0] * a[1][0] + b[1] * a[1][1] + b[2] * a[1][2], + b[0] * a[2][0] + b[1] * a[2][1] + b[2] * a[2][2], + ) + + @hjit("f(V,V)", inline=True) def matmul_VV_hf(a, b): return a[0] * b[0] + a[1] * b[1] + a[2] * b[2] From ab87ea8c24851016d79d8cd28bb6722644df9fe8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 19 Feb 2024 19:08:02 +0100 Subject: [PATCH 311/346] jit elevation function --- src/hapsira/core/events.py | 60 ++++++++++++++++------------- src/hapsira/twobody/orbit/scalar.py | 11 ++---- 2 files changed, 37 insertions(+), 34 deletions(-) diff --git a/src/hapsira/core/events.py b/src/hapsira/core/events.py index 069e2cf3e..eed0e1c65 100644 --- a/src/hapsira/core/events.py +++ b/src/hapsira/core/events.py @@ -1,11 +1,8 @@ -from math import acos, cos, sin - -from numba import njit as jit -import numpy as np +from math import acos, asin, cos, sin, sqrt from .elements import coe_rotation_matrix_hf, rv2coe_hf, RV2COE_TOL from .jit import array_to_V_hf, hjit, gjit -from .math.linalg import matmul_VV_hf, norm_V_hf +from .math.linalg import div_Vs_hf, matmul_MV_hf, matmul_VV_hf, norm_V_hf, sub_VV_hf from .util import planetocentric_to_AltAz_hf @@ -15,7 +12,8 @@ "eclipse_function_gf", "line_of_sight_hf", "line_of_sight_gf", - "elevation_function", + "elevation_function_hf", + "elevation_function_gf", ] @@ -135,17 +133,15 @@ def line_of_sight_gf(r1, r2, R, delta_theta): delta_theta[0] = line_of_sight_hf(array_to_V_hf(r1), array_to_V_hf(r2), R) -@jit -def elevation_function(k, u_, phi, theta, R, R_p, H): +@hjit("f(V,f,f,f,f,f)") +def elevation_function_hf(rr, phi, theta, R, R_p, H): """Calculates the elevation angle of an object in orbit with respect to a location on attractor. Parameters ---------- - k: float - Standard gravitational parameter. - u_: numpy.ndarray - Satellite position and velocity vector with respect to the central attractor. + rr: tuple[float,float,float] + Satellite position vector with respect to the central attractor. phi: float Geodetic Latitude of the station. theta: float @@ -157,26 +153,38 @@ def elevation_function(k, u_, phi, theta, R, R_p, H): H: float Elevation, above the ellipsoidal surface. """ - ecc = np.sqrt(1 - (R_p / R) ** 2) - denom = np.sqrt(1 - ecc**2 * np.sin(phi) ** 2) + + cos_phi = cos(phi) + sin_phi = sin(phi) + + ecc = sqrt(1 - (R_p / R) ** 2) + denom = sqrt(1 - ecc * ecc * sin_phi * sin_phi) g1 = H + (R / denom) - g2 = H + (1 - ecc**2) * R / denom + g2 = H + (1 - ecc * ecc) * R / denom + # Coordinates of location on attractor. - coords = np.array( - [ - g1 * np.cos(phi) * np.cos(theta), - g1 * np.cos(phi) * np.sin(theta), - g2 * np.sin(phi), - ] + coords = ( + g1 * cos_phi * cos(theta), + g1 * cos_phi * sin(theta), + g2 * sin_phi, ) # Position of satellite with respect to a point on attractor. - rho = np.subtract(u_[:3], coords) + rho = sub_VV_hf(rr, coords) - rot_matrix = np.array(planetocentric_to_AltAz_hf(theta, phi)) + rot_matrix = planetocentric_to_AltAz_hf(theta, phi) - new_rho = rot_matrix @ rho - new_rho = new_rho / np.linalg.norm(new_rho) - el = np.arcsin(new_rho[-1]) + new_rho = matmul_MV_hf(rot_matrix, rho) + new_rho = div_Vs_hf(new_rho, norm_V_hf(new_rho)) + el = asin(new_rho[-1]) return el + + +@gjit("void(f[:],f,f,f,f,f,f[:])", "(n),(),(),(),(),()->()") +def elevation_function_gf(rr, phi, theta, R, R_p, H, el): + """ + Vectorized elevation_function + """ + + el[0] = elevation_function_hf(array_to_V_hf(rr), phi, theta, R, R_p, H) diff --git a/src/hapsira/twobody/orbit/scalar.py b/src/hapsira/twobody/orbit/scalar.py index c7e27949a..fb4b115fe 100644 --- a/src/hapsira/twobody/orbit/scalar.py +++ b/src/hapsira/twobody/orbit/scalar.py @@ -11,7 +11,7 @@ import numpy as np from hapsira.bodies import Earth -from hapsira.core.events import elevation_function as elevation_function_fast +from hapsira.core.events import elevation_function_gf from hapsira.frames.util import get_frame from hapsira.threebody.soi import laplace_radius from hapsira.twobody.elements import eccentricity_vector, energy, t_p @@ -682,13 +682,8 @@ def elevation(self, lat, theta, h): "Elevation implementation is currently only supported for orbits having Earth as the attractor." ) - x, y, z = self.r.to_value(u.km) - vx, vy, vz = self.v.to_value(u.km / u.s) - u_ = np.array([x, y, z, vx, vy, vz]) - - elevation = elevation_function_fast( - self.attractor.k.to_value(u.km**3 / u.s**2), - u_, + elevation = elevation_function_gf( # pylint: disable=E1120 + self.r.to_value(u.km), lat.to_value(u.rad), theta.to_value(u.rad), self.attractor.R.to(u.km).value, From 984df4ed95ec09376e5cc5c4ea9eb086744e904a Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 23 Feb 2024 10:24:11 +0100 Subject: [PATCH 312/346] jit mean_motion --- src/hapsira/core/elements.py | 19 ++++++++++++++++++- src/hapsira/twobody/elements.py | 3 ++- src/hapsira/twobody/states.py | 12 ++++++++++-- 3 files changed, 30 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index f9a948451..e1600f24e 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -2,7 +2,7 @@ convert between different elements that define the orbit of a body. """ -from math import acos, atan, atan2, cos, log, pi, sin, sqrt, tan +from math import acos, atan, atan2, cos, fabs, log, pi, sin, sqrt, tan from .angles import E_to_nu_hf, F_to_nu_hf from .jit import array_to_V_hf, hjit, gjit, vjit @@ -38,6 +38,7 @@ "mee2coe_gf", "mee2rv_hf", "mee2rv_gf", + "mean_motion_vf", ] @@ -671,3 +672,19 @@ def mee2rv_gf(p, f, g, h, k, L, dummy, r, v): assert dummy.shape == (3,) (r[0], r[1], r[2]), (v[0], v[1], v[2]) = mee2rv_hf(p, f, g, h, k, L) + + +@hjit("f(f,f)", inline=True) +def mean_motion_hf(k, a): + """ + Mean motion given body (k) and semimajor axis (a). + """ + return sqrt(k / fabs(a * a * a)) + + +@vjit("f(f,f)") +def mean_motion_vf(k, a): + """ + Vectorized mean_motion + """ + return mean_motion_hf(k, a) diff --git a/src/hapsira/twobody/elements.py b/src/hapsira/twobody/elements.py index b6d05bd7a..5df0d297f 100644 --- a/src/hapsira/twobody/elements.py +++ b/src/hapsira/twobody/elements.py @@ -5,6 +5,7 @@ circular_velocity_vf, coe2rv_gf, eccentricity_vector_gf, + mean_motion_vf, ) from hapsira.core.propagation.farnocchia import delta_t_from_nu_vf, FARNOCCHIA_DELTA @@ -21,7 +22,7 @@ def circular_velocity(k, a): @u.quantity_input(k=u_km3s2, a=u.km) def mean_motion(k, a): """Mean motion given body (k) and semimajor axis (a).""" - return np.sqrt(k / abs(a**3)).to(1 / u.s) * u.rad + return mean_motion_vf(k.to_value(u_km3s2), a.to_value(u.km)) * u.rad / u.s @u.quantity_input(k=u_km3s2, a=u.km) diff --git a/src/hapsira/twobody/states.py b/src/hapsira/twobody/states.py index 1462cf4c1..22bcc1611 100644 --- a/src/hapsira/twobody/states.py +++ b/src/hapsira/twobody/states.py @@ -10,8 +10,9 @@ mee2rv_gf, rv2coe_gf, RV2COE_TOL, + mean_motion_vf, ) -from hapsira.twobody.elements import mean_motion, period, t_p +from hapsira.twobody.elements import period, t_p class BaseState: @@ -47,7 +48,14 @@ def attractor(self): @cached_property def n(self): """Mean motion.""" - return mean_motion(self.attractor.k, self.to_classical().a) + return ( + mean_motion_vf( + self.attractor.k.to_value(u.km**3 / u.s**2), + self.to_classical().a.to_value(u.km), + ) + * u.rad + / u.s + ) @cached_property def period(self): From 576a55ca5e2da45bd18936575fc89bda84093d5b Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 23 Feb 2024 10:37:10 +0100 Subject: [PATCH 313/346] jit period --- src/hapsira/core/elements.py | 19 +++++++++++++++++++ src/hapsira/twobody/elements.py | 4 ++-- src/hapsira/twobody/states.py | 16 +++++++++++++--- 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/src/hapsira/core/elements.py b/src/hapsira/core/elements.py index e1600f24e..151005e58 100644 --- a/src/hapsira/core/elements.py +++ b/src/hapsira/core/elements.py @@ -39,6 +39,7 @@ "mee2rv_hf", "mee2rv_gf", "mean_motion_vf", + "period_vf", ] @@ -688,3 +689,21 @@ def mean_motion_vf(k, a): Vectorized mean_motion """ return mean_motion_hf(k, a) + + +@hjit("f(f,f)", inline=True) +def period_hf(k, a): + """ + Period given body (k) and semimajor axis (a). + """ + n = mean_motion_hf(k, a) + return 2 * pi / n + + +@vjit("f(f,f)") +def period_vf(k, a): + """ + Vectorized period + """ + + return period_hf(k, a) diff --git a/src/hapsira/twobody/elements.py b/src/hapsira/twobody/elements.py index 5df0d297f..735c21e1f 100644 --- a/src/hapsira/twobody/elements.py +++ b/src/hapsira/twobody/elements.py @@ -6,6 +6,7 @@ coe2rv_gf, eccentricity_vector_gf, mean_motion_vf, + period_vf, ) from hapsira.core.propagation.farnocchia import delta_t_from_nu_vf, FARNOCCHIA_DELTA @@ -28,8 +29,7 @@ def mean_motion(k, a): @u.quantity_input(k=u_km3s2, a=u.km) def period(k, a): """Period given body (k) and semimajor axis (a).""" - n = mean_motion(k, a) - return 2 * np.pi * u.rad / n + return period_vf(k.to_value(u_km3s2), a.to_value(u.km)) * u.s @u.quantity_input(k=u_km3s2, r=u.km, v=u_kms) diff --git a/src/hapsira/twobody/states.py b/src/hapsira/twobody/states.py index 22bcc1611..c7b9ce71e 100644 --- a/src/hapsira/twobody/states.py +++ b/src/hapsira/twobody/states.py @@ -11,8 +11,12 @@ rv2coe_gf, RV2COE_TOL, mean_motion_vf, + period_vf, ) -from hapsira.twobody.elements import period, t_p +from hapsira.twobody.elements import t_p + + +u_km3s2 = u.km**3 / u.s**2 class BaseState: @@ -50,7 +54,7 @@ def n(self): """Mean motion.""" return ( mean_motion_vf( - self.attractor.k.to_value(u.km**3 / u.s**2), + self.attractor.k.to_value(u_km3s2), self.to_classical().a.to_value(u.km), ) * u.rad @@ -60,7 +64,13 @@ def n(self): @cached_property def period(self): """Period of the orbit.""" - return period(self.attractor.k, self.to_classical().a) + return ( + period_vf( + self.attractor.k.to_value(u_km3s2), + self.to_classical().a.to_value(u.km), + ) + * u.s + ) @cached_property def r_p(self): From 2b5752c2828caed09b134d6e066485a392b682a5 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 23 Feb 2024 10:51:09 +0100 Subject: [PATCH 314/346] jit t_p cleanup --- src/hapsira/twobody/states.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/hapsira/twobody/states.py b/src/hapsira/twobody/states.py index c7b9ce71e..f80e65e00 100644 --- a/src/hapsira/twobody/states.py +++ b/src/hapsira/twobody/states.py @@ -13,7 +13,7 @@ mean_motion_vf, period_vf, ) -from hapsira.twobody.elements import t_p +from hapsira.core.propagation.farnocchia import delta_t_from_nu_vf, FARNOCCHIA_DELTA u_km3s2 = u.km**3 / u.s**2 @@ -85,11 +85,16 @@ def r_a(self): @cached_property def t_p(self): """Elapsed time since latest perifocal passage.""" - return t_p( - self.to_classical().nu, - self.to_classical().ecc, - self.attractor.k, - self.r_p, + self_classical = self.to_classical() + return ( + delta_t_from_nu_vf( + self_classical.nu.to_value(u.rad), + self_classical.ecc.value, + self.attractor.k.to_value(u_km3s2), + self.r_p.to_value(u.km), + FARNOCCHIA_DELTA, + ) + * u.s ) def to_tuple(self): From 9793a8d91a68ca2b8aa7e4bfb3bad351ab9f8676 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 23 Feb 2024 13:21:07 +0100 Subject: [PATCH 315/346] fix core import contract --- pyproject.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 3389fb6c5..403d9199f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -177,6 +177,10 @@ name = "hapsira.core does not import astropy.units" type = "forbidden" source_modules = ["hapsira.core"] forbidden_modules = ["astropy.units"] +ignore_imports = [ + "hapsira.core.earth.atmosphere.coesa76 -> astropy.io.ascii", + "hapsira.core.earth.atmosphere.coesa76 -> astropy.utils.data" +] [tool.pytest.ini_options] norecursedirs = [".git", ".tox", "dist", "build", ".venv"] From eba70589ecb60218dac678bae59d435c3ac7a10c Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 23 Feb 2024 14:19:35 +0100 Subject: [PATCH 316/346] tox prevents numba cache, i.e. parallel compile causes segfaults -> run tests single-threaded on CI and new envs --- tox.ini | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tox.ini b/tox.ini index a62b0fc6a..61d32887b 100644 --- a/tox.ini +++ b/tox.ini @@ -25,10 +25,10 @@ setenv = PIP_PREFER_BINARY = 1 NPY_DISABLE_CPU_FEATURES = AVX512_SKX coverage: NUMBA_DISABLE_JIT = 1 - fast: PYTEST_MARKERS = -m "not slow and not mpl_image_compare" -n auto - online: PYTEST_MARKERS = -m "remote_data" -n auto - slow: PYTEST_MARKERS = -m "slow" -n auto - images: PYTEST_MARKERS = -m "mpl_image_compare" -n auto + fast: PYTEST_MARKERS = -m "not slow and not mpl_image_compare" -n 1 + online: PYTEST_MARKERS = -m "remote_data" -n 1 + slow: PYTEST_MARKERS = -m "slow" -n 1 + images: PYTEST_MARKERS = -m "mpl_image_compare" -n 1 PYTEST_EXTRA_ARGS = --mypy online: PYTEST_EXTRA_ARGS = --remote-data=any slow: PYTEST_EXTRA_ARGS = From 2d9e3b01fdbbde2ed1171b9633210557965399ad Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 23 Feb 2024 14:44:46 +0100 Subject: [PATCH 317/346] parallel CI without cache? --- tox.ini | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tox.ini b/tox.ini index 61d32887b..6c44bce33 100644 --- a/tox.ini +++ b/tox.ini @@ -24,11 +24,12 @@ setenv = PYTHONUNBUFFERED = yes PIP_PREFER_BINARY = 1 NPY_DISABLE_CPU_FEATURES = AVX512_SKX + HAPSIRA_CACHE = 0 coverage: NUMBA_DISABLE_JIT = 1 - fast: PYTEST_MARKERS = -m "not slow and not mpl_image_compare" -n 1 - online: PYTEST_MARKERS = -m "remote_data" -n 1 - slow: PYTEST_MARKERS = -m "slow" -n 1 - images: PYTEST_MARKERS = -m "mpl_image_compare" -n 1 + fast: PYTEST_MARKERS = -m "not slow and not mpl_image_compare" -n auto + online: PYTEST_MARKERS = -m "remote_data" -n auto + slow: PYTEST_MARKERS = -m "slow" -n auto + images: PYTEST_MARKERS = -m "mpl_image_compare" -n auto PYTEST_EXTRA_ARGS = --mypy online: PYTEST_EXTRA_ARGS = --remote-data=any slow: PYTEST_EXTRA_ARGS = From 1f3a910673dea46888c7a78cba0fcac13a0aef2e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 2 Mar 2024 17:57:45 +0100 Subject: [PATCH 318/346] fix plots --- .../examples/natural-and-artificial-perturbations.myst.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/examples/natural-and-artificial-perturbations.myst.md b/docs/source/examples/natural-and-artificial-perturbations.myst.md index 1d1f5bb3d..d87ffb21c 100644 --- a/docs/source/examples/natural-and-artificial-perturbations.myst.md +++ b/docs/source/examples/natural-and-artificial-perturbations.myst.md @@ -229,9 +229,9 @@ ephem = initial.to_ephem( ```{code-cell} ipython3 frame = OrbitPlotter(backend=Plotly3D()) - frame.set_attractor(Earth) frame.plot_ephem(ephem, label="orbit influenced by Moon") +frame.show() ``` ## Applying thrust @@ -282,9 +282,9 @@ ephem2 = orb0.to_ephem( ```{code-cell} ipython3 frame = OrbitPlotter(backend=Plotly3D()) - frame.set_attractor(Earth) frame.plot_ephem(ephem2, label="orbit with artificial thrust") +frame.show() ``` ## Combining multiple perturbations From bce3039976ffea53e7f350dc42d7ee20105a0e59 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 2 Mar 2024 18:16:10 +0100 Subject: [PATCH 319/346] fix urls --- ...ter-with-python-using-jupyter-and-poliastro.myst.md | 2 +- .../going-to-mars-with-python-using-poliastro.myst.md | 2 +- .../loading-OMM-and-TLE-satellite-data.myst.md | 4 ++-- src/hapsira/twobody/propagation/farnocchia.py | 2 +- src/hapsira/twobody/sampling.py | 4 ++-- tests/tests_plotting/test_orbit_plotter.py | 2 +- tests/tests_twobody/test_orbit.py | 10 +++++----- tests/tests_twobody/test_propagation.py | 6 +++--- 8 files changed, 16 insertions(+), 16 deletions(-) diff --git a/docs/source/examples/going-to-jupiter-with-python-using-jupyter-and-poliastro.myst.md b/docs/source/examples/going-to-jupiter-with-python-using-jupyter-and-poliastro.myst.md index 61c73e3b0..e77c29649 100644 --- a/docs/source/examples/going-to-jupiter-with-python-using-jupyter-and-poliastro.myst.md +++ b/docs/source/examples/going-to-jupiter-with-python-using-jupyter-and-poliastro.myst.md @@ -37,7 +37,7 @@ from hapsira.twobody import Orbit from hapsira.util import norm, time_range ``` -All the data for Juno's mission is sorted [here](https://github.com/hapsira/hapsira/wiki/EuroPython:-Per-Python-ad-Astra). The main maneuvers that the spacecraft will perform are listed down: +All the data for Juno's mission is sorted [here](https://github.com/poliastro/poliastro/wiki/EuroPython:-Per-Python-ad-Astra). The main maneuvers that the spacecraft will perform are listed down: * Inner cruise phase 1: This will set Juno in a new orbit around the sun. * Inner cruise phase 2: Fly-by around Earth. Gravity assist is performed. diff --git a/docs/source/examples/going-to-mars-with-python-using-poliastro.myst.md b/docs/source/examples/going-to-mars-with-python-using-poliastro.myst.md index 676c8ed7a..6c59378e0 100644 --- a/docs/source/examples/going-to-mars-with-python-using-poliastro.myst.md +++ b/docs/source/examples/going-to-mars-with-python-using-poliastro.myst.md @@ -13,7 +13,7 @@ kernelspec: # Going to Mars with Python using hapsira -This is an example on how to use [hapsira](https://github.com/hapsira/hapsira), a little library I've been working on to use in my Astrodynamics lessons. It features conversion between **classical orbital elements** and position vectors, propagation of **Keplerian orbits**, initial orbit determination using the solution of the **Lambert's problem** and **orbit plotting**. +This is an example on how to use [hapsira](https://github.com/pleiszenburg/hapsira), a little library I've been working on to use in my Astrodynamics lessons. It features conversion between **classical orbital elements** and position vectors, propagation of **Keplerian orbits**, initial orbit determination using the solution of the **Lambert's problem** and **orbit plotting**. In this example we're going to draw the trajectory of the mission [Mars Science Laboratory (MSL)](http://mars.jpl.nasa.gov/msl/), which carried the rover Curiosity to the surface of Mars in a period of something less than 9 months. diff --git a/docs/source/examples/loading-OMM-and-TLE-satellite-data.myst.md b/docs/source/examples/loading-OMM-and-TLE-satellite-data.myst.md index a4778698e..b4c9c3b97 100644 --- a/docs/source/examples/loading-OMM-and-TLE-satellite-data.myst.md +++ b/docs/source/examples/loading-OMM-and-TLE-satellite-data.myst.md @@ -28,7 +28,7 @@ kernelspec: +++ -However, it turns out that GP data in general, and TLEs in particular, are poorly understood even by professionals ([[1]](https://www.linkedin.com/posts/tom-johnson-32333a2_flawed-data-activity-6825845118990381056-yJX7), [[2]](https://twitter.com/flightclubio/status/1435303066085982209), [[3]](https://github.com/hapsira/hapsira/issues/1185)). The core issue is that TLEs and OMMs contain _Brouwer mean elements_, which **cannot be directly translated to osculating elements**. +However, it turns out that GP data in general, and TLEs in particular, are poorly understood even by professionals ([[1]](https://www.linkedin.com/posts/tom-johnson-32333a2_flawed-data-activity-6825845118990381056-yJX7), [[2]](https://twitter.com/flightclubio/status/1435303066085982209), [[3]](https://github.com/poliastro/poliastro/issues/1185)). The core issue is that TLEs and OMMs contain _Brouwer mean elements_, which **cannot be directly translated to osculating elements**. From "Spacetrack Report #3": @@ -58,7 +58,7 @@ Therefore, the **correct** way of using GP data is: As explained in the [Orbit Mean-Elements Messages (OMMs) support assessment](https://opensatcom.org/2020/12/28/omm-assessment-sgp4-benchmarks/) deliverable of OpenSatCom, OMM input/output support in open source libraries is somewhat scattered. Luckily, [python-sgp4](https://pypi.org/project/sgp4/) supports reading OMM in CSV and XML format, as well as usual TLE and 3LE formats. On the other hand, Astropy has accurate transformations from TEME to other reference frames. ```{code-cell} ipython3 -# From https://github.com/hapsira/hapsira/blob/main/contrib/satgpio.py +# From https://github.com/pleiszenburg/hapsira/blob/main/contrib/satgpio.py """ Author: Juan Luis Cano Rodríguez diff --git a/src/hapsira/twobody/propagation/farnocchia.py b/src/hapsira/twobody/propagation/farnocchia.py index d499a41fd..dbad0b863 100644 --- a/src/hapsira/twobody/propagation/farnocchia.py +++ b/src/hapsira/twobody/propagation/farnocchia.py @@ -52,7 +52,7 @@ def propagate_many(self, state, tofs): rv0 = state.to_value() # TODO: This should probably return a ClassicalStateArray instead, - # see discussion at https://github.com/hapsira/hapsira/pull/1492 + # see discussion at https://github.com/poliastro/poliastro/pull/1492 rr, vv = farnocchia_rv_gf(k, *rv0, tofs.to_value(u.s)) # pylint: disable=E0633 return ( diff --git a/src/hapsira/twobody/sampling.py b/src/hapsira/twobody/sampling.py index cd7a412f9..fde8eb20b 100644 --- a/src/hapsira/twobody/sampling.py +++ b/src/hapsira/twobody/sampling.py @@ -91,7 +91,7 @@ def sample(self, orbit): # However, we are also returning the epochs # (since computing them here is more efficient than doing it from the outside) # but there are open questions around StateArrays and epochs. - # See discussion at https://github.com/hapsira/hapsira/pull/1492 + # See discussion at https://github.com/poliastro/poliastro/pull/1492 cartesian = CartesianRepresentation( rr, differentials=CartesianDifferential(vv, xyz_axis=1), xyz_axis=1 ) @@ -160,7 +160,7 @@ def sample(self, orbit): # However, we are also returning the epochs # (since computing them here is more efficient than doing it from the outside) # but there are open questions around StateArrays and epochs. - # See discussion at https://github.com/hapsira/hapsira/pull/1492 + # See discussion at https://github.com/poliastro/poliastro/pull/1492 cartesian = CartesianRepresentation( rr, differentials=CartesianDifferential(vv, xyz_axis=1), xyz_axis=1 ) diff --git a/tests/tests_plotting/test_orbit_plotter.py b/tests/tests_plotting/test_orbit_plotter.py index b769979f8..897f21c62 100644 --- a/tests/tests_plotting/test_orbit_plotter.py +++ b/tests/tests_plotting/test_orbit_plotter.py @@ -242,7 +242,7 @@ def test_set_frame_plots_same_colors(): def test_redraw_keeps_trajectories(): - # See https://github.com/hapsira/hapsira/issues/518 + # See https://github.com/poliastro/poliastro/issues/518 op = OrbitPlotter() trajectory = churi.sample() op.plot_body_orbit(Mars, J2000_TDB, label="Mars") diff --git a/tests/tests_twobody/test_orbit.py b/tests/tests_twobody/test_orbit.py index 5eb5e63f7..59713ca33 100644 --- a/tests/tests_twobody/test_orbit.py +++ b/tests/tests_twobody/test_orbit.py @@ -419,7 +419,7 @@ def test_sample_numpoints(): def test_sample_big_orbits(): - # See https://github.com/hapsira/hapsira/issues/265 + # See https://github.com/poliastro/poliastro/issues/265 ss = Orbit.from_vectors( Sun, [-9_018_878.6, -94_116_055, 22_619_059] * u.km, @@ -1199,14 +1199,14 @@ def test_time_to_anomaly(expected_nu): # In some corner cases the resulting anomaly goes out of range, # and rather than trying to fix it right now # we will wait until we remove the round tripping, - # see https://github.com/hapsira/hapsira/issues/921 + # see https://github.com/poliastro/poliastro/issues/921 # FIXME: Add test that verifies that `orbit.nu` is always within range assert_quantity_allclose(iss_propagated.nu, expected_nu, atol=1e-12 * u.rad) @pytest.mark.xfail def test_can_set_iss_attractor_to_earth(): - # See https://github.com/hapsira/hapsira/issues/798 + # See https://github.com/poliastro/poliastro/issues/798 epoch = Time("2019-11-10 12:00:00") ephem = Ephem.from_horizons( "International Space Station", @@ -1235,7 +1235,7 @@ def test_issue_916(mock_query): def test_near_parabolic_M_does_not_hang(near_parabolic): - # See https://github.com/hapsira/hapsira/issues/907 + # See https://github.com/poliastro/poliastro/issues/907 expected_nu = -168.65 * u.deg orb = near_parabolic.propagate_to_anomaly(expected_nu) @@ -1253,7 +1253,7 @@ def test_propagation_near_parabolic_orbits_zero_seconds_gives_same_anomaly( def test_propagation_near_parabolic_orbits_does_not_hang(near_parabolic): - # See https://github.com/hapsira/hapsira/issues/475 + # See https://github.com/poliastro/poliastro/issues/475 orb_final = near_parabolic.propagate(near_parabolic.period) # Smoke test diff --git a/tests/tests_twobody/test_propagation.py b/tests/tests_twobody/test_propagation.py index 593e73db6..6de438993 100644 --- a/tests/tests_twobody/test_propagation.py +++ b/tests/tests_twobody/test_propagation.py @@ -347,7 +347,7 @@ def test_propagate_to_date_has_proper_epoch(): ) def test_propagate_long_times_keeps_geometry(method): # TODO: Extend to other propagators? - # See https://github.com/hapsira/hapsira/issues/265 + # See https://github.com/poliastro/poliastro/issues/265 time_of_flight = 100 * u.year res = iss.propagate(time_of_flight, method=method) @@ -436,7 +436,7 @@ def test_propagation_sets_proper_epoch(): def test_sample_around_moon_works(): - # See https://github.com/hapsira/hapsira/issues/649 + # See https://github.com/poliastro/poliastro/issues/649 orbit = Orbit.circular(Moon, 100 << u.km) coords = orbit.sample(10) @@ -446,7 +446,7 @@ def test_sample_around_moon_works(): def test_propagate_around_moon_works(): - # See https://github.com/hapsira/hapsira/issues/649 + # See https://github.com/poliastro/poliastro/issues/649 orbit = Orbit.circular(Moon, 100 << u.km) new_orbit = orbit.propagate(1 << u.h) From 89a3b6742c3ab3a0cb88f2237929b1c76d7b0271 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 6 Mar 2024 10:06:53 +0100 Subject: [PATCH 320/346] add changes --- docs/source/changelog.md | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/docs/source/changelog.md b/docs/source/changelog.md index 6a5502281..ccbc3d3a5 100644 --- a/docs/source/changelog.md +++ b/docs/source/changelog.md @@ -1,4 +1,22 @@ -# What's new +# Changes + +## hapsira 0.19.0 - 2024-XX-XX + +**CAUTION**: A number changes at least partially **BREAK BACKWARDS COMPATIBILITY** for certain use cases. + +This release features a significant refactoring of `core`, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) for details. All relevant `core` functions are now designed to work equally on CPUs and GPUs as either [universal functions](https://numba.readthedocs.io/en/stable/user/vectorize.html#the-vectorize-decorator) or [generalized universal functions](https://numba.readthedocs.io/en/stable/user/vectorize.html#the-guvectorize-decorator). As a "side-effect", all relevant `core` functions allow parallel operation with full [broadcasting semantics](https://numpy.org/doc/stable/user/basics.broadcasting.html). Their single-thread performance was also increased depending on use-case by around two orders of magnitude. All refactored **functions** in `core` were **renamed**, now carrying additional suffixes to indicate how they can or can not be invoked. + +Critical fix changing behaviour: The Loss of Signal (LOS) event would previously produce wrong results. + +- FEATURE: New `core.math` module, including fast replacements for many `numpy` and some `scipy` functions, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) +- FEATURE: New `core.jit` module, wrapping a number of `numba` functions to have a central place to apply settings, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) +- FEATURE: New `settings` module, mainly for handling JIT compiler settings, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) +- FEATURE: New `debug` module, including logging capabilities, mainly logging JIT compiler issues, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) +- FIX: The Loss of Signal (LOS) event would misshandle the position of the secondary body i.e. producing wrong results, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) as well as the [relevant commit](https://github.com/pleiszenburg/hapsira/commit/988a91cd22ff1de285c33af35b13d288963fcaf7) +- FIX: The Cowell propagator could produce wrong results if times of flight (tof) where provided in units other than seconds, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) +- FIX: Broken plots in example notebooks, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) +- FIX: Typo in `bodies`, see [hapsira #6](https://github.com/pleiszenburg/hapsira/pull/6) +- DEV: Parallel (multi-core) testing enabled by default, see [hapsira #5](https://github.com/pleiszenburg/hapsira/pull/5) ## hapsira 0.18.0 - 2023-12-24 From 009b718e4cc51c4a0475a5814ca77514af2bddbc Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 6 Mar 2024 10:21:38 +0100 Subject: [PATCH 321/346] fix copyright --- docs/source/conf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 244d5dae9..20fc4979f 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -87,7 +87,7 @@ } project = "hapsira" -copyright = "2023 Sebastian M. Ernst" +copyright = "2023-2024 Sebastian M. Ernst" project_ver = version(project) version = ".".join(project_ver.split(".")[:2]) From 17b6493dce315370040f6b3dd87b3e3ef65271b6 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 6 Mar 2024 14:55:36 +0100 Subject: [PATCH 322/346] details --- docs/source/changelog.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/source/changelog.md b/docs/source/changelog.md index ccbc3d3a5..79886e7f4 100644 --- a/docs/source/changelog.md +++ b/docs/source/changelog.md @@ -8,7 +8,9 @@ This release features a significant refactoring of `core`, see [hapsira #7](http Critical fix changing behaviour: The Loss of Signal (LOS) event would previously produce wrong results. -- FEATURE: New `core.math` module, including fast replacements for many `numpy` and some `scipy` functions, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) +- FEATURE: New `core.math` module, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7), including fast replacements for many `numpy` and some `scipy` functions, most notably: + - [scipy.interpolate.interp1d](https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.interp1d.html) is replaced by `core.math.interpolate.interp_hb`. It custom-compiles 1D linear interpolators, embedding data statically into the compiled functions. + - [scipy.integrate.solve_ivp](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html), [scipy.integrate.DOP853](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.DOP853.html) and [scipy.optimize.brentq](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.brentq.html) are replaced by `core.math.ivp`, a purely functional compiled implementation running entirely on the stack. - FEATURE: New `core.jit` module, wrapping a number of `numba` functions to have a central place to apply settings, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) - FEATURE: New `settings` module, mainly for handling JIT compiler settings, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) - FEATURE: New `debug` module, including logging capabilities, mainly logging JIT compiler issues, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) From b6dac8a345142e236f269c64873bc213eecea8ad Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Wed, 6 Mar 2024 21:26:31 +0100 Subject: [PATCH 323/346] draft core doc --- docs/source/core.md | 208 +++++++++++++++++++++++++++++++++++++++++++ docs/source/index.md | 1 + 2 files changed, 209 insertions(+) create mode 100644 docs/source/core.md diff --git a/docs/source/core.md b/docs/source/core.md new file mode 100644 index 000000000..95177455d --- /dev/null +++ b/docs/source/core.md @@ -0,0 +1,208 @@ +(quickstart)= +# The core module + +The `core` module handles most actual heavy computations. It is compiled via [numba](https://numba.readthedocs.io). For both working with functions from `core` directly and contributing to it, it is highly recommended to gain some basic understanding of how `numba` works. + +`core` is designed to work equally on CPUs and GPUs. (Most) exported `core` APIs interfacing with the rest of `hapsira` are either [universal functions](https://numba.readthedocs.io/en/stable/user/vectorize.html#the-vectorize-decorator) or [generalized universal functions](https://numba.readthedocs.io/en/stable/user/vectorize.html#the-guvectorize-decorator), which allow parallel operation with full [broadcasting semantics](https://numpy.org/doc/stable/user/basics.broadcasting.html). + +```{warning} +Some `core` functions have yet not been refactored into this shape and will soon follow the same approach. +``` + +## Compiler targets + +There are three compiler targets, which can be controlled through settings and/or environment variables: + +- `cpu`: Single-threaded on CPUs +- `parallel`: Parallelized by `numba` via [threading layers](https://numba.readthedocs.io/en/stable/user/threading-layer.html) on CPUs +- `cuda`: Parallelized by `numba` via CUDA on Nvidia GPUs + +All code of `core` will be compiled for one of the above listed targets. If multiple targets are supposed to be used simultaneously, this can only be achieved by multiple Python processes running in parallel. + +## Compiler decorators + +`core` offers the follwing JIT compiler decorators provided via `core.jit`: + +- `vjit`: Wraps `numba.vectorize`. Functions decorated by it carry the suffix `_vf`. +- `gjit`: Wraps `numba.guvectorize`. Functions decorated by it carry the suffix `_gf`. +- `hjit`: Wraps `numba.jit` or `numba.cuda.jit`, depending on compiler target. Functions decorated by it carry the suffix `_hf`. +- `djit`: Variation of `hjit` with fixed function signature for user-provided functions used by `Cowell` + +`core` functions dynamically generating (and compiling) functions within their scope carry `_hb`, `_vb` and `_gb` suffixes. + +Wrapping `numba` functions allows to centralize compiler options and target switching as well as to simplify typing. + +The decorators are applied in a **hierarchy**: + +- Functions decorated by either `vjit` and `gjit` serve as the **only** interface between regular uncompiled Python code and `core` +- Functions decorated by `vjit` and `hjit` only call functions decorated by `hjit` +- Functions decorated by `hjit` can only call each other. +- Functions decorated by `vjit`, `gjit` and `hjit`: + - are only allowed to depend on Python's standard library's [math module](https://docs.python.org/3/library/math.html), but not on [numpy](https://numpy.org/doc/stable/) - except for certain details like [enforcing floating point precision](https://numpy.org/doc/stable/user/basics.types.html) as provided by `core.math.ieee754` + - are fully typed, loosely following [numba semantics](https://numba.readthedocs.io/en/stable/reference/types.html) via shortcuts + +```{note} +The above mentioned "hierarchy" of decorators is imposed by CUDA-compatibility. While functions decorated by `numba.jit` (targets `cpu` and `parallel`) can be called from uncompiled Python code, functions decorated by `numba.cuda.jit` (target `cuda`) are considered [device functions](https://numba.readthedocs.io/en/stable/cuda/device-functions.html) and can not be called by uncompiled Python code directly. They are supposed to be called by CUDA-kernels (or other device functions) only (slightly simplifying the actual situation as implemented by `numba`). If the target is set to `cuda`, functions decorated by `numba.vectorize` and `numba.guvectorize` become CUDA kernels. +``` + +```{note} +Eliminating `numpy` as a dependency serves two purposes. While it also contributes to [CUDA-compatiblity](https://numba.readthedocs.io/en/stable/cuda/cudapysupported.html), it additionally makes the code significantly faster on CPUs. +``` + +```{warning} +As a result of name suffixes as of version `0.19.0`, many `core` module functions have been renamed making the package intentionally backwards-incompatible. Functions not yet using the new infrastructure can be recognized based on lack of suffix. Eventually all `core` functions will use this infrastructure and carry matching suffixes. +``` + +## Typing + +All functions decorated by `hjit`, `vjit` and `gjit` must by typed using [signatures similar to those of numba](https://numba.readthedocs.io/en/stable/reference/types.html). + +All compiled code enforces a single floating point precision level, which can be configured. The default is FP64 / double precision. For simplicity, the type shortcut is `f`, replacing `f2`, `f4` or `f8`. Additional infrastructure can be found in `core.math.ieee754`. Consider the following example: + +```python +from numba import vectorize +from hapsira.core.jit import vjit + +@vectorize("f8(f8)") +def foo(x): + return x ** 2 + +@vjit("f(f)") +def bar_vf(x): + return x ** 2 +``` + +3D vectors are expressed as tuples, type shortcut `V`, replacing `Tuple([f,f,f])`. Consider the following example: + +```python +from numba import njit +from hapsira.core.jit import hjit + +@njit("f8(Tuple([f8,f8,f8]))") +def foo(x): + return x[0] + x[1] + x[2] + +@hjit("f(V)") +def bar_hf(x): + return x[0] + x[1] + x[2] +``` + +Matrices are expressed as tuples of tuples, type shortcut `M`, replacing `Tuple([V,V,V])`. Consider the following example: + +```python +from numba import njit +from hapsira.core.jit import hjit + +@njit("f8(Tuple([Tuple([f8,f8,f8]),Tuple([f8,f8,f8]),Tuple([f8,f8,f8])]))") +def foo(x): + sum_ = 0 + for idx in range(3): + for jdx in range(3): + sum_ += x[idx][jdx] + return sum_ + +@hjit("f(M)") +def bar_hf(x): + sum_ = 0 + for idx in range(3): + for jdx in range(3): + sum_ += x[idx][jdx] + return sum_ +``` + +Function types use the shortcut `F`, replacing `FunctionType`. + +## Cowell’s formulation + +Cowell’s formulation is one of the few places where `core` is exposed directly to the user. + +### Two-body function + +In its most simple form, the `CowellPropagator` relies on a variation of `func_twobody_hf` as a parameter, a function compiled by `hjit`, which can technically be omitted: + +```python +from hapsira.core.propagation.base import func_twobody_hf +from hapsira.twobody.propagation import CowellPropagator + +prop = CowellPropagator(f=func_twobody_hf) +prop = CowellPropagator() # identical to the above +``` + +If perturbations are applied, however, `func_twobody_hf` needs to be altered. It is important that the new altered function is compiled via the `hjit` decorator and that is has the correct signature. To simplify the matter for users, a variation of `hjit` named `djit` carries the correct signature implicitly: + +```python +from hapsira.core.jit import djit, hjit +from hapsira.core.math.linalg import mul_Vs_hf +from hapsira.core.propagation.base import func_twobody_hf +from hapsira.twobody.propagation import CowellPropagator + +@hjit("Tuple([V,V])(f,V,V,f)") +def foo_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + return du_kep_rr, mul_Vs_hf(du_kep_vv, 1.1) # multiply speed vector by 1.1 + +@djit +def bar_hf(t0, rr, vv, k): + du_kep_rr, du_kep_vv = func_twobody_hf(t0, rr, vv, k) + return du_kep_rr, mul_Vs_hf(du_kep_vv, 1.1) # multiply speed vector by 1.1 + +prop = CowellPropagator(f=foo_hf) +prop = CowellPropagator(f=bar_hf) # identical to the above +``` + +### Events + +The core of each event's implementation must also be compiled by `hjit`. New events must inherit from `BaseEvent`. The compiled implementation should be an attribute, a function or a static method, named `_impl_hf`. Once this attribute is specified, an explicit call to the `_wrap` method, most likely from the constructor, automatically generates a second version of `_impl_hf` named `_impl_dense_hf` that is used to not only evaluate but also to approximate the exact time of flight of an event based on dense output of the underlying solver. + +## Settings + +The following settings, available via `settings.settings`, allow to alter the compiler's behaviour: + +- `DEBUG`: `bool`, default `False` +- `CACHE`: `bool`, default `not DEBUG` +- `TARGET`: `str`, default `cpu`, alternatives `parallel` and `cuda` +- `INLINE`: `bool`, default `TARGET == "cuda"` +- `NOPYTHON`: `bool`, default `True` +- `FORCEOBJ`: `bool`, default `False` +- `PRECISION`: `str`, default `f8`, alternatives `f2` and `f4` + +```{note} +Settings can be switched by either setting environment variables or importing the `settings` module **before** any other (sub-) module is imported. +``` + +The `DEBUG` setting disables caching and enables the highest log level, among other things. + +`CACHE` only works for `cpu` and `parallel` targets. It speeds up import times drastically if the package gets reused. Dynamically generated functions can not be cached and must be exempt from caching by passing `cache = False` as a parameter to the JIT compiler decorator. + +```{warning} +Building the cache should not be done in parallel processes - this will result in segmentation faults, see [numba #4807](https://github.com/numba/numba/issues/4807). Once `core` is fully compiled and cached, it can however be used in parallel processes. Rebuilding the cache can usually reliably resolve segmentation faults. +``` + +Inlining via `INLINE` drastically increases performance but also compile times. It is the default behaviour for target `cuda`. See [relevant chapter in numba documentation](https://numba.readthedocs.io/en/stable/developer/inlining.html#notes-on-inlining) for details. + +`NOPYTHON` and `FORCEOBJ` provide additional debugging capabilities but should not be changed for regular use. For details, see [nopython mode](https://numba.readthedocs.io/en/stable/glossary.html#term-nopython-mode) and [object mode](https://numba.readthedocs.io/en/stable/glossary.html#term-object-mode) in `numba`'s documentation. + +The default `PRECISION` of all floating point operations is FP64 / double precision float. + +```{warning} +`hapsira`, formerly `poliastro`, was validated using FP64. Certain parts like Cowell reliably operate at this precision only. Other parts like for instance atmospheric models can easily handle single precision. This option is therefore provided for experimental purposes only. +``` + +## Logging + +Compiler issues are logged via logging channel `hapsira` using Python's standard library's [logging module](https://docs.python.org/3/howto/logging.html), also available as `debug.logger`. All compiler activity can be observed by enabling log level `debug`. + +## Math + +The former `_math` module, version `0.18` and earlier, has become a first-class citizen as `core.math`, fully compiled by the above mentioned infrastructure. `core.math` contains a number of replacements for `numpy` operations, mostly found in `core.math.linalg`. All of those functions do not allocate memory and are free of side-effects including a lack of changes to their parameters. + +Functions in `core.math` follow a loose naming convention, indicating for what types of parameters they can be used. `mul_Vs_hf` for instance is a multiplication of a vector `V` and a scalar `s` (floating point). + +`core.math` also replaces (some) required `scipy` functions: + +- [scipy.interpolate.interp1d](https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.interp1d.html) is replaced by `core.math.interpolate.interp_hb`. It custom-compiles 1D linear interpolators, embedding data statically into the compiled functions. +- [scipy.integrate.solve_ivp](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html), [scipy.integrate.DOP853](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.DOP853.html) and [scipy.optimize.brentq](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.brentq.html) are replaced by `core.math.ivp`. + +```{note} +Future releases might remove more dependencies to `scipy` from `core` for full CUDA compatibility and additional performance. +``` diff --git a/docs/source/index.md b/docs/source/index.md index 581eba39d..c60ecd940 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -161,6 +161,7 @@ caption: How-to guides & Examples --- gallery contributing +core ``` ```{toctree} From 0a49ca971696949ae3dfc42a2f9c99658008a404 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 7 Mar 2024 11:16:28 +0100 Subject: [PATCH 324/346] clarifications --- docs/source/core.md | 23 ++++++++++++----------- 1 file changed, 12 insertions(+), 11 deletions(-) diff --git a/docs/source/core.md b/docs/source/core.md index 95177455d..416057519 100644 --- a/docs/source/core.md +++ b/docs/source/core.md @@ -37,22 +37,23 @@ The decorators are applied in a **hierarchy**: - Functions decorated by either `vjit` and `gjit` serve as the **only** interface between regular uncompiled Python code and `core` - Functions decorated by `vjit` and `hjit` only call functions decorated by `hjit` - Functions decorated by `hjit` can only call each other. -- Functions decorated by `vjit`, `gjit` and `hjit`: - - are only allowed to depend on Python's standard library's [math module](https://docs.python.org/3/library/math.html), but not on [numpy](https://numpy.org/doc/stable/) - except for certain details like [enforcing floating point precision](https://numpy.org/doc/stable/user/basics.types.html) as provided by `core.math.ieee754` - - are fully typed, loosely following [numba semantics](https://numba.readthedocs.io/en/stable/reference/types.html) via shortcuts ```{note} -The above mentioned "hierarchy" of decorators is imposed by CUDA-compatibility. While functions decorated by `numba.jit` (targets `cpu` and `parallel`) can be called from uncompiled Python code, functions decorated by `numba.cuda.jit` (target `cuda`) are considered [device functions](https://numba.readthedocs.io/en/stable/cuda/device-functions.html) and can not be called by uncompiled Python code directly. They are supposed to be called by CUDA-kernels (or other device functions) only (slightly simplifying the actual situation as implemented by `numba`). If the target is set to `cuda`, functions decorated by `numba.vectorize` and `numba.guvectorize` become CUDA kernels. -``` - -```{note} -Eliminating `numpy` as a dependency serves two purposes. While it also contributes to [CUDA-compatiblity](https://numba.readthedocs.io/en/stable/cuda/cudapysupported.html), it additionally makes the code significantly faster on CPUs. +The "hierarchy" of decorators is imposed by CUDA-compatibility. While functions decorated by `numba.jit` (targets `cpu` and `parallel`) can be called from uncompiled Python code, functions decorated by `numba.cuda.jit` (target `cuda`) are considered [device functions](https://numba.readthedocs.io/en/stable/cuda/device-functions.html) and can not be called by uncompiled Python code directly. They are supposed to be called by CUDA-kernels (or other device functions) only (slightly simplifying the actual situation as implemented by `numba`). If the target is set to `cuda`, functions decorated by `numba.vectorize` and `numba.guvectorize` become CUDA kernels. ``` ```{warning} As a result of name suffixes as of version `0.19.0`, many `core` module functions have been renamed making the package intentionally backwards-incompatible. Functions not yet using the new infrastructure can be recognized based on lack of suffix. Eventually all `core` functions will use this infrastructure and carry matching suffixes. ``` +## Dependencies + +Functions decorated by `vjit`, `gjit` and `hjit` are only allowed to depend on Python's standard library's [math module](https://docs.python.org/3/library/math.html), but **not** on other third-party packages like [numpy](https://numpy.org/doc/stable/) or [scipy](https://docs.scipy.org/doc/scipy/) for that matter - except for certain details like [enforcing floating point precision](https://numpy.org/doc/stable/user/basics.types.html) as provided by `core.math.ieee754` + +```{note} +Eliminating `numpy` and other dependencies serves two purposes. While it is critical for [CUDA-compatiblity](https://numba.readthedocs.io/en/stable/cuda/cudapysupported.html), it additionally makes the code significantly faster on CPUs. +``` + ## Typing All functions decorated by `hjit`, `vjit` and `gjit` must by typed using [signatures similar to those of numba](https://numba.readthedocs.io/en/stable/reference/types.html). @@ -175,7 +176,7 @@ The `DEBUG` setting disables caching and enables the highest log level, among ot `CACHE` only works for `cpu` and `parallel` targets. It speeds up import times drastically if the package gets reused. Dynamically generated functions can not be cached and must be exempt from caching by passing `cache = False` as a parameter to the JIT compiler decorator. ```{warning} -Building the cache should not be done in parallel processes - this will result in segmentation faults, see [numba #4807](https://github.com/numba/numba/issues/4807). Once `core` is fully compiled and cached, it can however be used in parallel processes. Rebuilding the cache can usually reliably resolve segmentation faults. +Building the cache should not be done in parallel processes - this will most likely result in non-deterministic segmentation faults, see [numba #4807](https://github.com/numba/numba/issues/4807). Once `core` is fully compiled and cached, it can however be used in parallel processes. Rebuilding the cache can usually reliably resolve segmentation faults. ``` Inlining via `INLINE` drastically increases performance but also compile times. It is the default behaviour for target `cuda`. See [relevant chapter in numba documentation](https://numba.readthedocs.io/en/stable/developer/inlining.html#notes-on-inlining) for details. @@ -185,7 +186,7 @@ Inlining via `INLINE` drastically increases performance but also compile times. The default `PRECISION` of all floating point operations is FP64 / double precision float. ```{warning} -`hapsira`, formerly `poliastro`, was validated using FP64. Certain parts like Cowell reliably operate at this precision only. Other parts like for instance atmospheric models can easily handle single precision. This option is therefore provided for experimental purposes only. +`hapsira`, formerly `poliastro`, was validated for FP64. Certain parts like Cowell reliably operate at this precision only. Other parts like for instance atmospheric models can easily handle single precision. This option is therefore provided for experimental purposes only. ``` ## Logging @@ -196,7 +197,7 @@ Compiler issues are logged via logging channel `hapsira` using Python's standard The former `_math` module, version `0.18` and earlier, has become a first-class citizen as `core.math`, fully compiled by the above mentioned infrastructure. `core.math` contains a number of replacements for `numpy` operations, mostly found in `core.math.linalg`. All of those functions do not allocate memory and are free of side-effects including a lack of changes to their parameters. -Functions in `core.math` follow a loose naming convention, indicating for what types of parameters they can be used. `mul_Vs_hf` for instance is a multiplication of a vector `V` and a scalar `s` (floating point). +Functions in `core.math` follow a loose naming convention, indicating for what types of parameters they can be used. `mul_Vs_hf` for instance is a multiplication of a vector `V` and a scalar `s` (floating point). `M` indicates matricis. `core.math` also replaces (some) required `scipy` functions: From 284604f4f1797689c96e03e009db4997034ebc34 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 7 Mar 2024 11:28:11 +0100 Subject: [PATCH 325/346] fix headline --- docs/source/core.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/core.md b/docs/source/core.md index 416057519..d261821fa 100644 --- a/docs/source/core.md +++ b/docs/source/core.md @@ -1,5 +1,5 @@ (quickstart)= -# The core module +# Core module The `core` module handles most actual heavy computations. It is compiled via [numba](https://numba.readthedocs.io). For both working with functions from `core` directly and contributing to it, it is highly recommended to gain some basic understanding of how `numba` works. From c131f30491f550bf1795bc88ffb8796466209fac Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 7 Mar 2024 11:30:23 +0100 Subject: [PATCH 326/346] cleanup --- src/hapsira/core/math/ivp/_const.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/hapsira/core/math/ivp/_const.py b/src/hapsira/core/math/ivp/_const.py index 69f54fedc..ff276ba04 100644 --- a/src/hapsira/core/math/ivp/_const.py +++ b/src/hapsira/core/math/ivp/_const.py @@ -21,7 +21,6 @@ MAX_FACTOR = 10 # Maximum allowed increase in a step size. INTERPOLATOR_POWER = 7 -N_STAGES_EXTENDED = 16 ERROR_ESTIMATOR_ORDER = 7 ERROR_EXPONENT = -1 / (ERROR_ESTIMATOR_ORDER + 1) From 3cf672497b756845d944701cb82d0e0c9e66400e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 7 Mar 2024 12:39:28 +0100 Subject: [PATCH 327/346] fix doc build issues --- docs/source/changelog.md | 2 +- docs/source/core.md | 2 +- ...ng-to-jupiter-with-python-using-jupyter-and-hapsira.myst.md} | 0 ....myst.md => going-to-mars-with-python-using-hapsira.myst.md} | 0 ...ps-with-poliastro.myst.md => porkchops-with-hapsira.myst.md} | 0 src/hapsira/core/spheroid_location.py | 2 +- src/hapsira/core/thrust/change_ecc_inc.py | 2 +- 7 files changed, 4 insertions(+), 4 deletions(-) rename docs/source/examples/{going-to-jupiter-with-python-using-jupyter-and-poliastro.myst.md => going-to-jupiter-with-python-using-jupyter-and-hapsira.myst.md} (100%) rename docs/source/examples/{going-to-mars-with-python-using-poliastro.myst.md => going-to-mars-with-python-using-hapsira.myst.md} (100%) rename docs/source/examples/{porkchops-with-poliastro.myst.md => porkchops-with-hapsira.myst.md} (100%) diff --git a/docs/source/changelog.md b/docs/source/changelog.md index 79886e7f4..be610ecdd 100644 --- a/docs/source/changelog.md +++ b/docs/source/changelog.md @@ -253,7 +253,7 @@ as well as the results from Google Summer of Code 2021. The interactive orbit plotters {py:class}`~poliastro.plotting.OrbitPlotter2D` and {py:class}`~poliastro.plotting.OrbitPlotter3D` now have a new method to easily display impulsive burns. - See {doc}`/examples/going-to-jupiter-with-python-using-jupyter-and-poliastro` + See {doc}`/examples/going-to-jupiter-with-python-using-jupyter-and-hapsira` for an example. - **Many performance improvements** Several contributors have helped accelerate more algorithms diff --git a/docs/source/core.md b/docs/source/core.md index d261821fa..93262eeaa 100644 --- a/docs/source/core.md +++ b/docs/source/core.md @@ -1,4 +1,4 @@ -(quickstart)= +(coremodule)= # Core module The `core` module handles most actual heavy computations. It is compiled via [numba](https://numba.readthedocs.io). For both working with functions from `core` directly and contributing to it, it is highly recommended to gain some basic understanding of how `numba` works. diff --git a/docs/source/examples/going-to-jupiter-with-python-using-jupyter-and-poliastro.myst.md b/docs/source/examples/going-to-jupiter-with-python-using-jupyter-and-hapsira.myst.md similarity index 100% rename from docs/source/examples/going-to-jupiter-with-python-using-jupyter-and-poliastro.myst.md rename to docs/source/examples/going-to-jupiter-with-python-using-jupyter-and-hapsira.myst.md diff --git a/docs/source/examples/going-to-mars-with-python-using-poliastro.myst.md b/docs/source/examples/going-to-mars-with-python-using-hapsira.myst.md similarity index 100% rename from docs/source/examples/going-to-mars-with-python-using-poliastro.myst.md rename to docs/source/examples/going-to-mars-with-python-using-hapsira.myst.md diff --git a/docs/source/examples/porkchops-with-poliastro.myst.md b/docs/source/examples/porkchops-with-hapsira.myst.md similarity index 100% rename from docs/source/examples/porkchops-with-poliastro.myst.md rename to docs/source/examples/porkchops-with-hapsira.myst.md diff --git a/src/hapsira/core/spheroid_location.py b/src/hapsira/core/spheroid_location.py index ef21f347e..6c552f329 100644 --- a/src/hapsira/core/spheroid_location.py +++ b/src/hapsira/core/spheroid_location.py @@ -26,7 +26,7 @@ def cartesian_cords(a, c, lon, lat, h): """Calculates cartesian coordinates. - Parametersnorm_V_hf + Parameters ---------- a : float Semi-major axis diff --git a/src/hapsira/core/thrust/change_ecc_inc.py b/src/hapsira/core/thrust/change_ecc_inc.py index 9a71aac91..4b2ebc2ac 100644 --- a/src/hapsira/core/thrust/change_ecc_inc.py +++ b/src/hapsira/core/thrust/change_ecc_inc.py @@ -3,7 +3,7 @@ References ---------- * Pollard, J. E. "Simplified Analysis of Low-Thrust Orbital Maneuvers", 2000. -rv2coe + """ from math import asin, atan, cos, pi, log, sin From 13a9994d659a4c5061d1286d5fd8703de1d81187 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 7 Mar 2024 17:43:47 +0100 Subject: [PATCH 328/346] rm type S entirely; inline array to vector --- src/hapsira/core/jit.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index 6b7875303..b9dc11675 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -71,10 +71,6 @@ def _parse_signatures(signature: str, noreturn: bool = False) -> Union[str, List signature = signature.replace("M", "Tuple([V,V,V])") # matrix is a tuple of vectors signature = signature.replace("V", "Tuple([f,f,f])") # vector is a tuple of floats - # signature = signature.replace( - # "S", "Tuple([f,f,f,f,f,f])" - # ) # state, two vectors, is a tuple of floats TODO remove, use vectors instead? - assert "S" not in signature signature = signature.replace( "F", "FunctionType" ) # TODO does not work for CUDA yet @@ -322,6 +318,6 @@ def wrapper(inner_func: Callable) -> Callable: return wrapper -@hjit("V(f[:])") +@hjit("V(f[:])", inline=True) def array_to_V_hf(x): return x[0], x[1], x[2] From cbc1acd3f19d87fcc572cd9ea5a5fc2591df85de Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 7 Mar 2024 17:45:09 +0100 Subject: [PATCH 329/346] more changes --- docs/source/changelog.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/source/changelog.md b/docs/source/changelog.md index be610ecdd..e9671ec4c 100644 --- a/docs/source/changelog.md +++ b/docs/source/changelog.md @@ -8,17 +8,24 @@ This release features a significant refactoring of `core`, see [hapsira #7](http Critical fix changing behaviour: The Loss of Signal (LOS) event would previously produce wrong results. +Module layout change: `core.earth_atmosphere` became `core.earth.atmosphere`. + - FEATURE: New `core.math` module, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7), including fast replacements for many `numpy` and some `scipy` functions, most notably: - [scipy.interpolate.interp1d](https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.interp1d.html) is replaced by `core.math.interpolate.interp_hb`. It custom-compiles 1D linear interpolators, embedding data statically into the compiled functions. - [scipy.integrate.solve_ivp](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.solve_ivp.html), [scipy.integrate.DOP853](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.DOP853.html) and [scipy.optimize.brentq](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.brentq.html) are replaced by `core.math.ivp`, a purely functional compiled implementation running entirely on the stack. - FEATURE: New `core.jit` module, wrapping a number of `numba` functions to have a central place to apply settings, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) - FEATURE: New `settings` module, mainly for handling JIT compiler settings, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) - FEATURE: New `debug` module, including logging capabilities, mainly logging JIT compiler issues, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) +- FEATURE: Significant portions of the COESA76 atmophere model are now compiled and available as part of `core`, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7). `core.earth_atmosphere` was renamed into `core.earth.atmosphere`. +- DOCS: The `core` module is technically user-facing and as such now explicitly documented, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) +- DOCS: The "quickstart" section received an update and now includes all previously missing imports. - FIX: The Loss of Signal (LOS) event would misshandle the position of the secondary body i.e. producing wrong results, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) as well as the [relevant commit](https://github.com/pleiszenburg/hapsira/commit/988a91cd22ff1de285c33af35b13d288963fcaf7) - FIX: The Cowell propagator could produce wrong results if times of flight (tof) where provided in units other than seconds, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) - FIX: Broken plots in example notebooks, see [hapsira #7](https://github.com/pleiszenburg/hapsira/pull/7) - FIX: Typo in `bodies`, see [hapsira #6](https://github.com/pleiszenburg/hapsira/pull/6) +- FIX: Some notebooks in the documentation had disappeared due to incomplete rebranding - DEV: Parallel (multi-core) testing enabled by default, see [hapsira #5](https://github.com/pleiszenburg/hapsira/pull/5) +- DEV: Deactivated warning due to too many simultaneously opened `matplotlib` plots ## hapsira 0.18.0 - 2023-12-24 From 72b757223875eef5194ab1e66bf5908ba7c55935 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 7 Mar 2024 17:45:30 +0100 Subject: [PATCH 330/346] ieee754; errors; kwargs --- docs/source/core.md | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/docs/source/core.md b/docs/source/core.md index 93262eeaa..4efb54707 100644 --- a/docs/source/core.md +++ b/docs/source/core.md @@ -46,6 +46,18 @@ The "hierarchy" of decorators is imposed by CUDA-compatibility. While functions As a result of name suffixes as of version `0.19.0`, many `core` module functions have been renamed making the package intentionally backwards-incompatible. Functions not yet using the new infrastructure can be recognized based on lack of suffix. Eventually all `core` functions will use this infrastructure and carry matching suffixes. ``` +```{note} +Some functions decorated by `gjit` must receive a dummy parameter, also explicitly named `dummy`. It is usually an empty `numpy` array of shape `(3,)` of data type `u1` (unsigned one-byte integer). This is a work-around for [numba #2797](https://github.com/numba/numba/issues/2797). +``` + +## Compiler errors + +Misconfigured compiler decorators or unavailable targets raise an `errors.JitError` exception. + +## Keyword arguments and defaults + +Due to incompletely documented limitations in `numba`, see [documentation](https://numba.readthedocs.io/en/stable/reference/pysupported.html#function-calls) and [numba #7870](https://github.com/numba/numba/issues/7870), functions decorated by `hjit`, `vjit` and `gjit` can not have defaults for any of their arguments. In this context, those functions can not reliably be called with keyword arguments, too, which must therefore be avoided. Defaults are provided as constants within the same submodule, usually the function name in capital letters followed by the name of the argument, also in capital letters. + ## Dependencies Functions decorated by `vjit`, `gjit` and `hjit` are only allowed to depend on Python's standard library's [math module](https://docs.python.org/3/library/math.html), but **not** on other third-party packages like [numpy](https://numpy.org/doc/stable/) or [scipy](https://docs.scipy.org/doc/scipy/) for that matter - except for certain details like [enforcing floating point precision](https://numpy.org/doc/stable/user/basics.types.html) as provided by `core.math.ieee754` @@ -58,7 +70,7 @@ Eliminating `numpy` and other dependencies serves two purposes. While it is crit All functions decorated by `hjit`, `vjit` and `gjit` must by typed using [signatures similar to those of numba](https://numba.readthedocs.io/en/stable/reference/types.html). -All compiled code enforces a single floating point precision level, which can be configured. The default is FP64 / double precision. For simplicity, the type shortcut is `f`, replacing `f2`, `f4` or `f8`. Additional infrastructure can be found in `core.math.ieee754`. Consider the following example: +All compiled code enforces a single floating point precision level, which can be configured. The default is FP64 / double precision. For simplicity, the type shortcut is `f`, replacing `f2`, `f4` or `f8`. Consider the following example: ```python from numba import vectorize @@ -73,6 +85,12 @@ def bar_vf(x): return x ** 2 ``` +Additional infrastructure can be found in `core.math.ieee754`. The default floating point type is exposed as `core.math.ieee754.float_` for explicit conversions. A matching epsilon is exposed as `core.math.ieee754.EPS`. + +```{note} +Divisions by zero should, regardless of compiler target or even entirely deactivated compiler, always result in `inf` (infinity) instead of `ZeroDivisionError` exceptions. Most divisions within `core` are therefore explicitly guarded. +``` + 3D vectors are expressed as tuples, type shortcut `V`, replacing `Tuple([f,f,f])`. Consider the following example: ```python From 82794a8bb7aea9f5ddda8c431db858ae14b8cf51 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 7 Mar 2024 20:17:25 +0100 Subject: [PATCH 331/346] inlining broke CI --- src/hapsira/core/jit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/jit.py b/src/hapsira/core/jit.py index b9dc11675..51d97f7cd 100644 --- a/src/hapsira/core/jit.py +++ b/src/hapsira/core/jit.py @@ -318,6 +318,6 @@ def wrapper(inner_func: Callable) -> Callable: return wrapper -@hjit("V(f[:])", inline=True) +@hjit("V(f[:])") def array_to_V_hf(x): return x[0], x[1], x[2] From 2dfc5f2e10629a9efc083023b347c2134ff08ed6 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Apr 2024 17:30:54 +0200 Subject: [PATCH 332/346] fix imports --- .../examples/natural-and-artificial-perturbations.myst.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/examples/natural-and-artificial-perturbations.myst.md b/docs/source/examples/natural-and-artificial-perturbations.myst.md index d87ffb21c..0a64d80a3 100644 --- a/docs/source/examples/natural-and-artificial-perturbations.myst.md +++ b/docs/source/examples/natural-and-artificial-perturbations.myst.md @@ -25,7 +25,7 @@ from hapsira.bodies import Earth, Moon from hapsira.constants import rho0_earth, H0_earth from hapsira.core.elements import rv2coe_gf, RV2COE_TOL -from hapsira.core.jit import array_to_V_hf, djit, hjit +from hapsira.core.jit import djit, hjit from hapsira.core.math.linalg import add_VV_hf from hapsira.core.perturbations import ( atmospheric_drag_exponential_hf, From f04baffffff0371924d6aad6bc17185fa3d6638d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sun, 21 Apr 2024 17:31:26 +0200 Subject: [PATCH 333/346] simplify exception for numba --- src/hapsira/core/propagation/cowell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hapsira/core/propagation/cowell.py b/src/hapsira/core/propagation/cowell.py index 9a4c8ebdd..28b85def7 100644 --- a/src/hapsira/core/propagation/cowell.py +++ b/src/hapsira/core/propagation/cowell.py @@ -216,7 +216,7 @@ def cowell_gf( event_g_olds[event_idx] = event_g_news[event_idx] if not t_last <= t: - raise ValueError("not t_last <= t", t_last, t) + raise ValueError("not t_last <= t") while t_idx[0] < tofs.shape[0] and tofs[t_idx[0]] < t: rrs[t_idx[0], :], vvs[t_idx[0], :] = dop853_dense_interp_hf( From 2a57e7b710078a55c74a6fb04ddc96640bdb089d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 22 Apr 2024 15:58:19 +0200 Subject: [PATCH 334/346] fix doc strings --- src/hapsira/core/earth/atmosphere/coesa.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/hapsira/core/earth/atmosphere/coesa.py b/src/hapsira/core/earth/atmosphere/coesa.py index 5c8d60315..8435d4706 100644 --- a/src/hapsira/core/earth/atmosphere/coesa.py +++ b/src/hapsira/core/earth/atmosphere/coesa.py @@ -21,7 +21,7 @@ def _geometric_to_geopotential_hf(z, r0): Returns ------- - h: float + h : float Geopotential altitude. """ h = r0 * z / (r0 + z) @@ -41,7 +41,7 @@ def _geopotential_to_geometric_hf(h, r0): Returns ------- - z: float + z : float Geometric altitude. """ z = r0 * h / (r0 - h) @@ -67,7 +67,7 @@ def _gravity_hf(z, g0, r0): Returns ------- - g: float + g : float Gravity value at given geometric altitude. """ g = g0 * (r0 / (r0 + z)) ** 2 @@ -87,7 +87,7 @@ def get_index_hf(x, x_levels): Returns ------- - i: int + i : int Index for the value. """ @@ -102,7 +102,7 @@ def get_index_hf(x, x_levels): @hjit("Tuple([f,f])(f,f,b1)") def check_altitude_hf(alt, r0, geometric): - # Get geometric and geopotential altitudes + """Get geometric and geopotential altitudes""" if geometric: z = alt h = _z_to_h_hf(z, r0) From 61afbe37fbcbb1ce68d09d4eb4c892f88f2c6da3 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 22 Apr 2024 16:00:24 +0200 Subject: [PATCH 335/346] fix doc strings --- src/hapsira/core/earth/atmosphere/coesa76.py | 106 ++++++++++++++----- 1 file changed, 77 insertions(+), 29 deletions(-) diff --git a/src/hapsira/core/earth/atmosphere/coesa76.py b/src/hapsira/core/earth/atmosphere/coesa76.py index 4c1c9752d..555c19c28 100644 --- a/src/hapsira/core/earth/atmosphere/coesa76.py +++ b/src/hapsira/core/earth/atmosphere/coesa76.py @@ -94,7 +94,7 @@ @hjit("Tuple([f,f])(f,f,b1)") -def _check_altitude_hf(alt, r0, geometric): # geometric True by default +def _check_altitude_hf(alt, r0, geometric): """Checks if altitude is inside valid range. Parameters @@ -105,16 +105,16 @@ def _check_altitude_hf(alt, r0, geometric): # geometric True by default Attractor radius. geometric : bool If `True`, assumes geometric altitude kind. + Default `True`. Returns ------- - z: float + z : float Geometric altitude. - h: float + h : float Geopotential altitude. """ - z, h = check_altitude_hf(alt, r0, geometric) assert zb_levels[0] <= z <= zb_levels[-1] @@ -132,8 +132,8 @@ def _get_index_zb_levels_hf(x): Returns ------- - i: int - Index for the value. + i : int + Index for the value. `999` if there was an error. """ for i, value in enumerate(zb_levels): @@ -156,8 +156,9 @@ def _get_index_z_coeff_hf(x): Returns ------- - i: int + i : int Index for the value. + Index for the value. `999` if there was an error. """ for i, value in enumerate(z_coeff): @@ -180,8 +181,9 @@ def _get_coefficients_avobe_86_p_coeff_hf(z): Returns ------- - coeffs: tuple[float,float,float,float,float] - List of corresponding coefficients + coeffs : tuple[float,float,float,float,float] + Tuple of corresponding coefficients + """ # Get corresponding coefficients i = _get_index_z_coeff_hf(z) @@ -199,8 +201,9 @@ def _get_coefficients_avobe_86_rho_coeff_hf(z): Returns ------- - coeffs: tuple[float,float,float,float,float] - List of corresponding coefficients + coeffs : tuple[float,float,float,float,float] + Tuple of corresponding coefficients + """ # Get corresponding coefficients i = _get_index_z_coeff_hf(z) @@ -214,7 +217,7 @@ def _get_coefficients_avobe_86_rho_coeff_hf(z): @hjit("f(f,b1)") -def temperature_hf(alt, geometric): # geometric True by default +def temperature_hf(alt, geometric): """Solves for temperature at given altitude. Parameters @@ -223,11 +226,13 @@ def temperature_hf(alt, geometric): # geometric True by default Geometric/Geopotential altitude. geometric : bool If `True`, assumes geometric altitude kind. + Default `True`. Returns ------- - T: float + T : float Kinetic temeperature. + """ # Test if altitude is inside valid range z, h = _check_altitude_hf(alt, r0, geometric) @@ -266,16 +271,29 @@ def temperature_hf(alt, geometric): # geometric True by default @vjit("f(f,b1)") -def temperature_vf(alt, geometric): # geometric True by default - """ - Vectorized temperature - """ +def temperature_vf(alt, geometric): + """Solves for temperature at given altitude. + Vectorized. + Parameters + ---------- + alt : float + Geometric/Geopotential altitude. + geometric : bool + If `True`, assumes geometric altitude kind. + Default `True`. + + Returns + ------- + T : float + Kinetic temeperature. + + """ return temperature_hf(alt, geometric) @hjit("f(f,b1)") -def pressure_hf(alt, geometric): # geometric True by default +def pressure_hf(alt, geometric): """Solves pressure at given altitude. Parameters @@ -284,11 +302,13 @@ def pressure_hf(alt, geometric): # geometric True by default Geometric/Geopotential altitude. geometric : bool If `True`, assumes geometric altitude kind. + Default `True`. Returns ------- - p: float + p : float Pressure at given altitude. + """ # Test if altitude is inside valid range z, h = _check_altitude_hf(alt, r0, geometric) @@ -322,16 +342,29 @@ def pressure_hf(alt, geometric): # geometric True by default @vjit("f(f,b1)") -def pressure_vf(alt, geometric): # geometric True by default - """ - Vectorized pressure - """ +def pressure_vf(alt, geometric): + """Solves pressure at given altitude. + Vectorized. + Parameters + ---------- + alt : float + Geometric/Geopotential altitude. + geometric : bool + If `True`, assumes geometric altitude kind. + Default `True`. + + Returns + ------- + p : float + Pressure at given altitude. + + """ return pressure_hf(alt, geometric) @hjit("f(f,b1)") -def density_hf(alt, geometric): # geometric True by default +def density_hf(alt, geometric): """Solves density at given height. Parameters @@ -340,11 +373,13 @@ def density_hf(alt, geometric): # geometric True by default Geometric/Geopotential height. geometric : bool If `True`, assumes that `alt` argument is geometric kind. + Default `True`. Returns ------- - rho: float + rho : float Density at given height. + """ # Test if altitude is inside valid range z, _ = _check_altitude_hf(alt, r0, geometric) @@ -368,9 +403,22 @@ def density_hf(alt, geometric): # geometric True by default @vjit("f(f,b1)") -def density_vf(alt, geometric): # geometric True by default - """ - Vectorized density - """ +def density_vf(alt, geometric): + """Solves density at given height. + Vectorized. + Parameters + ---------- + alt : float + Geometric/Geopotential height. + geometric : bool + If `True`, assumes that `alt` argument is geometric kind. + Default `True`. + + Returns + ------- + rho : float + Density at given height. + + """ return density_hf(alt, geometric) From 4a914d4514859538dea2aff3f83facd229517e5d Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Apr 2024 19:28:45 +0200 Subject: [PATCH 336/346] cleanup --- src/hapsira/core/earth/atmosphere/coesa76.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/hapsira/core/earth/atmosphere/coesa76.py b/src/hapsira/core/earth/atmosphere/coesa76.py index 555c19c28..43a02025a 100644 --- a/src/hapsira/core/earth/atmosphere/coesa76.py +++ b/src/hapsira/core/earth/atmosphere/coesa76.py @@ -273,7 +273,6 @@ def temperature_hf(alt, geometric): @vjit("f(f,b1)") def temperature_vf(alt, geometric): """Solves for temperature at given altitude. - Vectorized. Parameters ---------- @@ -344,7 +343,6 @@ def pressure_hf(alt, geometric): @vjit("f(f,b1)") def pressure_vf(alt, geometric): """Solves pressure at given altitude. - Vectorized. Parameters ---------- @@ -405,7 +403,6 @@ def density_hf(alt, geometric): @vjit("f(f,b1)") def density_vf(alt, geometric): """Solves density at given height. - Vectorized. Parameters ---------- From 0750ff457dd742d0c5a513f674ac43dbd86f0fda Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Apr 2024 19:38:12 +0200 Subject: [PATCH 337/346] doc strings --- src/hapsira/core/math/interpolate.py | 46 ++++++++++++++++++++++++++-- 1 file changed, 43 insertions(+), 3 deletions(-) diff --git a/src/hapsira/core/math/interpolate.py b/src/hapsira/core/math/interpolate.py index 780951ad8..4b46a299e 100644 --- a/src/hapsira/core/math/interpolate.py +++ b/src/hapsira/core/math/interpolate.py @@ -15,7 +15,22 @@ def interp_hb(x: np.ndarray, y: np.ndarray) -> Callable: """ - Build compiled 1d-interpolator for 3D vectors + Builds compiled linear 1D interpolator for 3D vectors, + embedding x and y as const values into the binary. + Does not extrapolate! + + Parameters + ---------- + x : np.ndarray + Values for x + y : np.ndarray + Values for y + + Returns + ------- + rho : Callable + 1D interpolator + """ assert x.ndim == 1 @@ -30,6 +45,21 @@ def interp_hb(x: np.ndarray, y: np.ndarray) -> Callable: @hjit("V(f)", cache=False) def interp_hf(x_new): + """ + 1D interpolator + + Parameters + ---------- + x_new : float + New value for x + + Returns + ------- + y_new : float + New value for y + + """ + assert x_new >= x[0] assert x_new <= x[-1] @@ -67,13 +97,20 @@ def interp_hf(x_new): def spline_interp(y, x, u, *, kind="cubic"): - """Interpolates y, sampled at x instants, at u instants using `scipy.interpolate.interp1d`.""" + """ + Interpolates y, sampled at x instants, at u instants using `scipy.interpolate.interp1d`. + + TODO compile + + """ + y_u = _scipy_interp1d(x, y, kind=kind)(u) return y_u def sinc_interp(y, x, u): - """Interpolates y, sampled at x instants, at u instants using sinc interpolation. + """ + Interpolates y, sampled at x instants, at u instants using sinc interpolation. Notes ----- @@ -82,7 +119,10 @@ def sinc_interp(y, x, u): see https://mail.python.org/pipermail/scipy-user/2012-January/031255.html. However, quick experiments show different ringing behavior. + TODO compile + """ + if len(y) != len(x): raise ValueError("x and s must be the same length") From c1790aad9c3220996f97215079368d28e5aa1fe2 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Apr 2024 20:25:37 +0200 Subject: [PATCH 338/346] doc strings; param name cleanup --- src/hapsira/core/math/linalg.py | 397 ++++++++++++++++++++++++++++++-- 1 file changed, 376 insertions(+), 21 deletions(-) diff --git a/src/hapsira/core/math/linalg.py b/src/hapsira/core/math/linalg.py index 7b0a419fa..503740cb4 100644 --- a/src/hapsira/core/math/linalg.py +++ b/src/hapsira/core/math/linalg.py @@ -27,22 +27,92 @@ @hjit("V(V)", inline=True) -def abs_V_hf(x): - return fabs(x[0]), fabs(x[1]), fabs(x[2]) +def abs_V_hf(a): + """ + Abs 3D vector of 3D vector element-wise. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + + Returns + ------- + b : tuple[float,float,float] + Vector + + """ + + return fabs(a[0]), fabs(a[1]), fabs(a[2]) @hjit("V(V,f)", inline=True) def add_Vs_hf(a, b): + """ + Adds a 3D vector and a scalar element-wise. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : float + Scalar + + Returns + ------- + c : tuple[float,float,float] + Vector + + """ + return a[0] + b, a[1] + b, a[2] + b @hjit("V(V,V)", inline=True) def add_VV_hf(a, b): + """ + Adds two 3D vectors. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : tuple[float,float,float] + Vector + + Returns + ------- + c : tuple[float,float,float] + Vector + + """ + return a[0] + b[0], a[1] + b[1], a[2] + b[2] @hjit("V(V,V)", inline=True) def cross_VV_hf(a, b): + """ + Cross-product of two 3D vectors. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : tuple[float,float,float] + Vector + + Returns + ------- + c : tuple[float,float,float] + Vector + + """ + return ( a[1] * b[2] - a[2] * b[1], a[2] * b[0] - a[0] * b[2], @@ -53,25 +123,98 @@ def cross_VV_hf(a, b): @hjit("f(f,f)", inline=True) def div_ss_hf(a, b): """ - Similar to np.divide + Division of two scalars. Similar to `numpy.divide` as it returns + +/- (depending on the sign of `a`) infinity if `b` is zero. + Required for compatibility if `core` is not compiled for debugging purposes. + Inline-compiled by default. + + Parameters + ---------- + a : float + Scalar + b : float + Scalar + + Returns + ------- + c : float + Scalar + """ + if b == 0: return inf if a >= 0 else -inf return a / b @hjit("V(V,V)", inline=True) -def div_VV_hf(x, y): - return div_ss_hf(x[0], y[0]), div_ss_hf(x[1], y[1]), div_ss_hf(x[2], y[2]) +def div_VV_hf(a, b): + """ + Division of two 3D vectors element-wise. Similar to `numpy.divide` as + it returns +/- (depending on the sign of `a`) infinity if `b` is zero. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : tuple[float,float,float] + Vector + + Returns + ------- + c : tuple[float,float,float] + Vector + + """ + + return div_ss_hf(a[0], b[0]), div_ss_hf(a[1], b[1]), div_ss_hf(a[2], b[2]) @hjit("V(V,f)", inline=True) -def div_Vs_hf(v, s): - return div_ss_hf(v[0], s), div_ss_hf(v[1], s), div_ss_hf(v[2], s) +def div_Vs_hf(a, b): + """ + Division of a 3D vector by a scalar element-wise. Similar to `numpy.divide` as + it returns +/- (depending on the sign of `a`) infinity if `b` is zero. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : float + Scalar + + Returns + ------- + c : tuple[float,float,float] + Vector + + """ + + return div_ss_hf(a[0], b), div_ss_hf(a[1], b), div_ss_hf(a[2], b) @hjit("M(M,M)", inline=True) def matmul_MM_hf(a, b): + """ + Matmul (dot product) between two 3x3 matrices. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[tuple[float,float,float],tuple[float,float,float],tuple[float,float,float]] + Matrix + b : tuple[tuple[float,float,float],tuple[float,float,float],tuple[float,float,float]] + Matrix + + Returns + ------- + c : tuple[tuple[float,float,float],tuple[float,float,float],tuple[float,float,float]] + Matrix + + """ + return ( ( a[0][0] * b[0][0] + a[0][1] * b[1][0] + a[0][2] * b[2][0], @@ -93,6 +236,25 @@ def matmul_MM_hf(a, b): @hjit("V(V,M)", inline=True) def matmul_VM_hf(a, b): + """ + Matmul (dot product) between a 3D row vector and a 3x3 matrix + resulting in a 3D vector. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : tuple[tuple[float,float,float],tuple[float,float,float],tuple[float,float,float]] + Matrix + + Returns + ------- + c : tuple[float,float,float] + Vector + + """ + return ( a[0] * b[0][0] + a[1] * b[1][0] + a[2] * b[2][0], a[0] * b[0][1] + a[1] * b[1][1] + a[2] * b[2][1], @@ -102,6 +264,25 @@ def matmul_VM_hf(a, b): @hjit("V(M,V)", inline=True) def matmul_MV_hf(a, b): + """ + Matmul (dot product) between a 3x3 matrix and a 3D column vector + resulting in a 3D vector. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[tuple[float,float,float],tuple[float,float,float],tuple[float,float,float]] + Matrix + b : tuple[float,float,float] + Vector + + Returns + ------- + c : tuple[float,float,float] + Vector + + """ + return ( b[0] * a[0][0] + b[1] * a[0][1] + b[2] * a[0][2], b[0] * a[1][0] + b[1] * a[1][1] + b[2] * a[1][2], @@ -111,60 +292,234 @@ def matmul_MV_hf(a, b): @hjit("f(V,V)", inline=True) def matmul_VV_hf(a, b): + """ + Matmul (dot product) between two 3D vectors resulting in a scalar. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : tuple[float,float,float] + Vector + + Returns + ------- + c : float + Scalar + + """ + return a[0] * b[0] + a[1] * b[1] + a[2] * b[2] @hjit("V(V,V)", inline=True) -def max_VV_hf(x, y): +def max_VV_hf(a, b): + """ + Max elements element-wise from two 3D vectors. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : tuple[float,float,float] + Vector + + Returns + ------- + c : tuple[float,float,float] + Vector + + """ + return ( - x[0] if x[0] > y[0] else y[0], - x[1] if x[1] > y[1] else y[1], - x[2] if x[2] > y[2] else y[2], + a[0] if a[0] > b[0] else b[0], + a[1] if a[1] > b[1] else b[1], + a[2] if a[2] > b[2] else b[2], ) @hjit("V(V,f)", inline=True) -def mul_Vs_hf(v, s): - return v[0] * s, v[1] * s, v[2] * s +def mul_Vs_hf(a, b): + """ + Multiplication of a 3D vector by a scalar element-wise. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : float + Scalar + + Returns + ------- + c : tuple[float,float,float] + Vector + + """ + + return a[0] * b, a[1] * b, a[2] * b @hjit("V(V,V)", inline=True) def mul_VV_hf(a, b): + """ + Multiplication of two 3D vectors element-wise. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : tuple[float,float,float] + Vector + + Returns + ------- + c : tuple[float,float,float] + Vector + + """ + return a[0] * b[0], a[1] * b[1], a[2] * b[2] @hjit("f(V)", inline=True) def norm_V_hf(a): + """ + Norm of a 3D vector. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + + Returns + ------- + b : float + Scalar + + """ + return sqrt(matmul_VV_hf(a, a)) @vjit("f(f,f,f)") def norm_V_vf(a, b, c): - # TODO add axis setting in some way for util.norm? + """ + Norm of a 3D vector. + + Parameters + ---------- + a : float + First dimension scalar + b : float + Second dimension scalar + c : float + Third dimension scalar + + Returns + ------- + d : float + Scalar + + """ + return norm_V_hf((a, b, c)) @hjit("f(V,V)", inline=True) -def norm_VV_hf(x, y): - return sqrt(matmul_VV_hf(x, x) + matmul_VV_hf(y, y)) +def norm_VV_hf(a, b): + """ + Combined norm of two 3D vectors treated like a single 6D vector. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : tuple[float,float,float] + Vector + + Returns + ------- + c : float + Scalar + + """ + + return sqrt(matmul_VV_hf(a, a) + matmul_VV_hf(b, b)) @hjit("f(f)", inline=True) -def sign_hf(x): - if x < 0.0: +def sign_hf(a): + """ + Sign of a float represented as another float (-1, 0, +1). + Inline-compiled by default. + + Parameters + ---------- + a : float + Scalar + + Returns + ------- + b : float + Scalar + + """ + + if a < 0.0: return -1.0 - if x == 0.0: + if a == 0.0: return 0.0 return 1.0 # if x > 0 @hjit("V(V,V)", inline=True) -def sub_VV_hf(va, vb): - return va[0] - vb[0], va[1] - vb[1], va[2] - vb[2] +def sub_VV_hf(a, b): + """ + Subtraction of two 3D vectors element-wise. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[float,float,float] + Vector + b : tuple[float,float,float] + Vector + + Returns + ------- + c : tuple[float,float,float] + Vector + + """ + + return a[0] - b[0], a[1] - b[1], a[2] - b[2] @hjit("M(M)", inline=True) def transpose_M_hf(a): + """ + Transposition of a matrix. + Inline-compiled by default. + + Parameters + ---------- + a : tuple[tuple[float,float,float],tuple[float,float,float],tuple[float,float,float]] + Matrix + + Returns + ------- + b : tuple[tuple[float,float,float],tuple[float,float,float],tuple[float,float,float]] + Matrix + + """ + return ( (a[0][0], a[1][0], a[2][0]), (a[0][1], a[1][1], a[2][1]), From 545c3f0a9dcc0c089d19d9aed71117b62bb7e829 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Apr 2024 20:33:46 +0200 Subject: [PATCH 339/346] doc strings --- src/hapsira/core/math/special.py | 105 ++++++++++++++++++++++++++++--- 1 file changed, 97 insertions(+), 8 deletions(-) diff --git a/src/hapsira/core/math/special.py b/src/hapsira/core/math/special.py index 34f142181..8c021630f 100644 --- a/src/hapsira/core/math/special.py +++ b/src/hapsira/core/math/special.py @@ -14,7 +14,8 @@ @hjit("f(f)") def hyp2f1b_hf(x): - """Hypergeometric function 2F1(3, 1, 5/2, x), see [Battin]. + """ + Hypergeometric function 2F1(3, 1, 5/2, x), see [Battin]. .. todo:: Add more information about this function @@ -24,7 +25,18 @@ def hyp2f1b_hf(x): More information about hypergeometric function can be checked at https://en.wikipedia.org/wiki/Hypergeometric_function + Parameters + ---------- + x : float + Scalar + + Returns + ------- + res : float + Scalar + """ + if x >= 1.0: return inf @@ -43,7 +55,26 @@ def hyp2f1b_hf(x): @vjit("f(f)") def hyp2f1b_vf(x): """ - Vectorized hyp2f1b + Hypergeometric function 2F1(3, 1, 5/2, x), see [Battin]. + + .. todo:: + Add more information about this function + + Notes + ----- + More information about hypergeometric function can be checked at + https://en.wikipedia.org/wiki/Hypergeometric_function + + Parameters + ---------- + x : float + Scalar + + Returns + ------- + b : float + Scalar + """ return hyp2f1b_hf(x) @@ -51,7 +82,8 @@ def hyp2f1b_vf(x): @hjit("f(f)") def stumpff_c2_hf(psi): - r"""Second Stumpff function. + r""" + Second Stumpff function. For positive arguments: @@ -59,7 +91,18 @@ def stumpff_c2_hf(psi): c_2(\psi) = \frac{1 - \cos{\sqrt{\psi}}}{\psi} + Parameters + ---------- + psi : float + Scalar + + Returns + ------- + res : float + Scalar + """ + eps = 1.0 if psi > eps: @@ -80,8 +123,25 @@ def stumpff_c2_hf(psi): @vjit("f(f)") def stumpff_c2_vf(psi): - """ - Vectorized stumpff_c2 + r""" + Second Stumpff function. + + For positive arguments: + + .. math:: + + c_2(\psi) = \frac{1 - \cos{\sqrt{\psi}}}{\psi} + + Parameters + ---------- + psi : float + Scalar + + Returns + ------- + res : float + Scalar + """ return stumpff_c2_hf(psi) @@ -89,7 +149,8 @@ def stumpff_c2_vf(psi): @hjit("f(f)") def stumpff_c3_hf(psi): - r"""Third Stumpff function. + r""" + Third Stumpff function. For positive arguments: @@ -97,7 +158,18 @@ def stumpff_c3_hf(psi): c_3(\psi) = \frac{\sqrt{\psi} - \sin{\sqrt{\psi}}}{\sqrt{\psi^3}} + Parameters + ---------- + psi : float + Scalar + + Returns + ------- + res : float + Scalar + """ + eps = 1.0 if psi > eps: @@ -118,8 +190,25 @@ def stumpff_c3_hf(psi): @vjit("f(f)") def stumpff_c3_vf(psi): - """ - Vectorized stumpff_c3 + r""" + Third Stumpff function. + + For positive arguments: + + .. math:: + + c_3(\psi) = \frac{\sqrt{\psi} - \sin{\sqrt{\psi}}}{\sqrt{\psi^3}} + + Parameters + ---------- + psi : float + Scalar + + Returns + ------- + res : float + Scalar + """ return stumpff_c3_hf(psi) From 292423748582d679fd9a5393ebb2167816658942 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Apr 2024 21:02:55 +0200 Subject: [PATCH 340/346] doc strings --- src/hapsira/core/math/ivp/_solve.py | 40 ++++++++++++++++++++++++----- 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/src/hapsira/core/math/ivp/_solve.py b/src/hapsira/core/math/ivp/_solve.py index ad6baf421..96b13a6e4 100644 --- a/src/hapsira/core/math/ivp/_solve.py +++ b/src/hapsira/core/math/ivp/_solve.py @@ -29,7 +29,25 @@ def dispatcher_hb( ) -> Callable: """ Workaround for https://github.com/numba/numba/issues/9420 + Compiles a dispatcher for a list of functions that can eventually called by index. + + Parameters + ---------- + funcs : tuple[Callable, ...] + One or multiple callables that require dispatching. + Dispatching will be based on position in tuple. + All callables must have the same signature. + argtypes : argument portion of signature for callables + restype : return type portion of signature for callables + arguments : names of arguments for callables + + Returns + ------- + b : Callable + Dispatcher function + """ + funcs = [ (f"func_{id(func):x}", func) for func in funcs ] # names are not unique, ids are @@ -60,20 +78,28 @@ def switch(idx): @hjit("b1(f,f,f)") def event_is_active_hf(g_old, g_new, direction): - """Find which event occurred during an integration step. + """ + Find which event occurred during an integration step. + + Based on + https://github.com/scipy/scipy/blob/4edfcaa3ce8a387450b6efce968572def71be089/scipy/integrate/_ivp/ivp.py#L130 Parameters ---------- - g, g_new : array_like, shape (n_events,) - Values of event functions at a current and next points. - directions : ndarray, shape (n_events,) - Event "direction" according to the definition in `solve_ivp`. + g_old : float + Value of event function at current point. + g_new : float + Value of event function at next point. + direction : float + Event "direction". Returns ------- - active_events : ndarray - Indices of events which occurred during the step. + active : boolean + Status of event (active or not) + """ + up = (g_old <= 0) & (g_new >= 0) down = (g_old >= 0) & (g_new <= 0) either = up | down From 2b16b5d984d39303e52cdbeceaa02f8ee8d677d6 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Apr 2024 21:33:17 +0200 Subject: [PATCH 341/346] doc strings --- src/hapsira/core/math/ivp/_brentq.py | 179 ++++++++++++++++++++++++++- 1 file changed, 175 insertions(+), 4 deletions(-) diff --git a/src/hapsira/core/math/ivp/_brentq.py b/src/hapsira/core/math/ivp/_brentq.py index f21d621c6..3371c3519 100644 --- a/src/hapsira/core/math/ivp/_brentq.py +++ b/src/hapsira/core/math/ivp/_brentq.py @@ -29,13 +29,47 @@ BRENTQ_MAXITER = 100 -@hjit("f(f,f)") +@hjit("f(f,f)", inline=True) def _min_ss_hf(a, b): + """ + The smaller of two scalars. + Inline by default. + + Parameters + ---------- + a : float + Scalar + b : float + Scalar + + Returns + ------- + c : float + Scalar + + """ + return a if a < b else b -@hjit("b1(f)") +@hjit("b1(f)", inline=True) def _signbit_s_hf(a): + """ + Sign bit of float. + Inline by default. + + Parameters + ---------- + a : float + Scalar + + Returns + ------- + b : boolean + Scalar + + """ + return a < 0 @@ -49,8 +83,45 @@ def _brentq_hf( maxiter, # int ): """ + Find a root of a function in a bracketing interval using Brent's method. + Loosely adapted from - https://github.com/scipy/scipy/blob/d23363809572e9a44074a3f06f66137083446b48/scipy/optimize/_zeros_py.py#L682 + - https://github.com/scipy/scipy/blob/d23363809572e9a44074a3f06f66137083446b48/scipy/optimize/_zeros_py.py#L682 + - https://github.com/scipy/scipy/blob/bd60b4ef9d886a9171345fc064c80aad6d171e73/scipy/optimize/Zeros/brentq.c#L37 + + Parameters + ---------- + func : Callable, float of float + The function :math:`f` must be continuous, and :math:`f(xa)` + and :math:`f(xb)` must have opposite signs. + xa : float + One end of the bracketing interval :math:`[xa, xb]`. + xb : float + The other end of the bracketing interval :math:`[xa, xb]`. + xtol : float + The computed root ``x0`` will satisfy ``np.allclose(x, x0, + atol=xtol, rtol=rtol)``, where ``x`` is the exact root. The + parameter must be positive. For nice functions, Brent's + method will often satisfy the above condition with ``xtol/2`` + and ``rtol/2``. + rtol : float + The computed root ``x0`` will satisfy ``np.allclose(x, x0, + atol=xtol, rtol=rtol)``, where ``x`` is the exact root. The + parameter cannot be smaller than its default value of + ``4*np.finfo(float).eps``. For nice functions, Brent's + method will often satisfy the above condition with ``xtol/2`` + and ``rtol/2``. + maxiter : float + If convergence is not achieved in `maxiter` iterations, an error is + raised. Must be >= 0. + + Returns + ------- + xc : float + Root of `f` between `a` and `b`. + status : int + Solver status code. + """ # if not xtol + 0. > 0: @@ -136,6 +207,14 @@ def _brentq_hf( def brentq_gb(func: Callable) -> Callable: """ Builds vectorized brentq + + Parameters + ---------- + func : Callable + + Returns + ------- + brentq_gf : Callable """ @gjit( @@ -152,6 +231,45 @@ def brentq_gf( xcur, status, ): + """ + Find a root of a function in a bracketing interval using Brent's method. + + Loosely adapted from + - https://github.com/scipy/scipy/blob/d23363809572e9a44074a3f06f66137083446b48/scipy/optimize/_zeros_py.py#L682 + - https://github.com/scipy/scipy/blob/bd60b4ef9d886a9171345fc064c80aad6d171e73/scipy/optimize/Zeros/brentq.c#L37 + + Parameters + ---------- + xa : float + One end of the bracketing interval :math:`[xa, xb]`. + xb : float + The other end of the bracketing interval :math:`[xa, xb]`. + xtol : float + The computed root ``x0`` will satisfy ``np.allclose(x, x0, + atol=xtol, rtol=rtol)``, where ``x`` is the exact root. The + parameter must be positive. For nice functions, Brent's + method will often satisfy the above condition with ``xtol/2`` + and ``rtol/2``. + rtol : float + The computed root ``x0`` will satisfy ``np.allclose(x, x0, + atol=xtol, rtol=rtol)``, where ``x`` is the exact root. The + parameter cannot be smaller than its default value of + ``4*np.finfo(float).eps``. For nice functions, Brent's + method will often satisfy the above condition with ``xtol/2`` + and ``rtol/2``. + maxiter : float + If convergence is not achieved in `maxiter` iterations, an error is + raised. Must be >= 0. + + Returns + ------- + xc : float + Root of `f` between `a` and `b`. + status : int + Solver status code. + + """ + xcur[0], status[0] = _brentq_hf(func, xa, xb, xtol, rtol, maxiter) return brentq_gf @@ -174,8 +292,61 @@ def brentq_dense_hf( argk, ): """ + Find a root of a function in a bracketing interval using Brent's method. + Virtually identical to `_brentq_hf`, except that it passes extra arguments + through to `func`. Architecturally required due to limiations of `numba`. + Loosely adapted from - https://github.com/scipy/scipy/blob/d23363809572e9a44074a3f06f66137083446b48/scipy/optimize/_zeros_py.py#L682 + - https://github.com/scipy/scipy/blob/d23363809572e9a44074a3f06f66137083446b48/scipy/optimize/_zeros_py.py#L682 + - https://github.com/scipy/scipy/blob/bd60b4ef9d886a9171345fc064c80aad6d171e73/scipy/optimize/Zeros/brentq.c#L37 + + Parameters + ---------- + func : Callable, float of float + The function :math:`f` must be continuous, and :math:`f(xa)` + and :math:`f(xb)` must have opposite signs. + idx : int + Selects function in dispatcher, passed to `func`. + xa : float + One end of the bracketing interval :math:`[xa, xb]`. + xb : float + The other end of the bracketing interval :math:`[xa, xb]`. + xtol : float + The computed root ``x0`` will satisfy ``np.allclose(x, x0, + atol=xtol, rtol=rtol)``, where ``x`` is the exact root. The + parameter must be positive. For nice functions, Brent's + method will often satisfy the above condition with ``xtol/2`` + and ``rtol/2``. + rtol : float + The computed root ``x0`` will satisfy ``np.allclose(x, x0, + atol=xtol, rtol=rtol)``, where ``x`` is the exact root. The + parameter cannot be smaller than its default value of + ``4*np.finfo(float).eps``. For nice functions, Brent's + method will often satisfy the above condition with ``xtol/2`` + and ``rtol/2``. + maxiter : float + If convergence is not achieved in `maxiter` iterations, an error is + raised. Must be >= 0. + sol1 : any + Passed to `func`. + sol2 : any + Passed to `func`. + sol3 : any + Passed to `func`. + sol4 : any + Passed to `func`. + sol5 : any + Passed to `func`. + argk : float + Passed to `func`. + + Returns + ------- + xc : float + Root of `f` between `a` and `b`. + status : int + Solver status code. + """ if not xtol + 0.0 > 0: From d9eac1ea77583aea39339095cc8388d3760b86e8 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Thu, 25 Apr 2024 21:35:05 +0200 Subject: [PATCH 342/346] scipy ref --- src/hapsira/core/math/ivp/_dop853_coefficients.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/hapsira/core/math/ivp/_dop853_coefficients.py b/src/hapsira/core/math/ivp/_dop853_coefficients.py index a425fc8ff..fec0bea2c 100644 --- a/src/hapsira/core/math/ivp/_dop853_coefficients.py +++ b/src/hapsira/core/math/ivp/_dop853_coefficients.py @@ -11,6 +11,9 @@ "E5", ] +# Based on +# https://github.com/scipy/scipy/blob/v1.12.0/scipy/integrate/_ivp/dop853_coefficients.py + C = np.array( [ 0.0, From e4cba2eb2343937d8021d331c008cb44c3794d64 Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Fri, 26 Apr 2024 23:35:55 +0200 Subject: [PATCH 343/346] prep doc --- src/hapsira/core/math/ivp/_rkcore.py | 58 ++++++++++++++++++++++++---- 1 file changed, 51 insertions(+), 7 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkcore.py b/src/hapsira/core/math/ivp/_rkcore.py index 12ceb83eb..573b1db94 100644 --- a/src/hapsira/core/math/ivp/_rkcore.py +++ b/src/hapsira/core/math/ivp/_rkcore.py @@ -54,6 +54,51 @@ def dop853_init_hf(fun, t0, rr, vv, t_bound, argk, rtol, atol): """ Explicit Runge-Kutta method of order 8. + Functional re-write of constructor of class `DOP853` within `scipy.integrate`. + + Based on + https://github.com/scipy/scipy/blob/4edfcaa3ce8a387450b6efce968572def71be089/scipy/integrate/_ivp/rk.py#L502 + + Parameters + ---------- + fun : float + Scalar + t0 : float + Scalar + rr : float + Scalar + vv : float + Scalar + t_bound : float + Scalar + argk : float + Scalar + rtol : float + Scalar + atol : float + Scalar + + Returns + ------- + t0 + rr + vv + t_bound + fun, # 4 + argk, # 5 + rtol, # 6 + atol, # 7 + direction, # 8 + K, # 9 + rr_old, # 10 + vv_old, # 11 + t_old, # 12 + h_previous, # 13 + status, # 14 + fr, # 15 + fv, # 16 + h_abs, # 17 + """ assert atol >= 0 @@ -150,14 +195,13 @@ def dop853_step_hf( fv, h_abs, ): - """Perform one integration step. + """ + Perform one integration step. + Functional re-write of method `step` of class `OdeSolver` within `scipy.integrate`. + + Based on + https://github.com/scipy/scipy/blob/4edfcaa3ce8a387450b6efce968572def71be089/scipy/integrate/_ivp/base.py#L175 - Returns - ------- - message : string or None - Report from the solver. Typically a reason for a failure if - `self.status` is 'failed' after the step was taken or None - otherwise. """ if status != DOP853_RUNNING: From 2c5fbc4d9bc5937fd7977bf381ba9cdc96cc7c5c Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Sat, 27 Apr 2024 00:07:23 +0200 Subject: [PATCH 344/346] doc strings --- src/hapsira/core/math/ivp/_rkcore.py | 152 ++++++++++++++++++++++----- 1 file changed, 125 insertions(+), 27 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkcore.py b/src/hapsira/core/math/ivp/_rkcore.py index 573b1db94..c70667348 100644 --- a/src/hapsira/core/math/ivp/_rkcore.py +++ b/src/hapsira/core/math/ivp/_rkcore.py @@ -57,47 +57,67 @@ def dop853_init_hf(fun, t0, rr, vv, t_bound, argk, rtol, atol): Functional re-write of constructor of class `DOP853` within `scipy.integrate`. Based on - https://github.com/scipy/scipy/blob/4edfcaa3ce8a387450b6efce968572def71be089/scipy/integrate/_ivp/rk.py#L502 + - https://github.com/scipy/scipy/blob/4edfcaa3ce8a387450b6efce968572def71be089/scipy/integrate/_ivp/rk.py#L502 + - https://github.com/scipy/scipy/blob/4edfcaa3ce8a387450b6efce968572def71be089/scipy/integrate/_ivp/rk.py#L85 + - https://github.com/scipy/scipy/blob/4edfcaa3ce8a387450b6efce968572def71be089/scipy/integrate/_ivp/base.py#L131 Parameters ---------- fun : float - Scalar + Right-hand side of the system. t0 : float - Scalar + Initial time. rr : float - Scalar + Initial state 0:3 vv : float - Scalar + Initial state 3:6 t_bound : float - Scalar + Boundary time argk : float - Scalar + Standard gravitational parameter for `fun` rtol : float - Scalar + Relative tolerance atol : float - Scalar + Absolute tolerance Returns ------- - t0 - rr - vv - t_bound - fun, # 4 - argk, # 5 - rtol, # 6 - atol, # 7 - direction, # 8 - K, # 9 - rr_old, # 10 - vv_old, # 11 - t_old, # 12 - h_previous, # 13 - status, # 14 - fr, # 15 - fv, # 16 - h_abs, # 17 + t0 : float + Initial time. + rr : tuple[float,float,float] + Initial state 0:3 + vv : tuple[float,float,float] + Initial state 3:6 + t_bound : float + Boundary time + fun : Callable + Right-hand side of the system + argk : float + Standard gravitational parameter for `fun` + rtol : float + Relative tolerance + atol : float + Absolute tolerance + direction : float + Integration direction + K : tuple[[float,...],...] + Storage array for RK stages + rr_old : tuple[float,float,float] + Last state 0:3 + vv_old : tuple[float,float,float] + Last state 3:6 + t_old : float + Last time + h_previous : float + Last step length + status : float + Solver status + fr : tuple[float,float,float] + Current value of the derivative 0:3 + fv : tuple[float,float,float] + Current value of the derivative 3:6 + h_abs : float + Absolute step """ @@ -202,6 +222,84 @@ def dop853_step_hf( Based on https://github.com/scipy/scipy/blob/4edfcaa3ce8a387450b6efce968572def71be089/scipy/integrate/_ivp/base.py#L175 + Parameters + ---------- + t : float + Current time. + rr : tuple[float,float,float] + Current state 0:3 + vv : tuple[float,float,float] + Current state 3:6 + t_bound : float + Boundary time + fun : Callable + Right-hand side of the system + argk : float + Standard gravitational parameter for `fun` + rtol : float + Relative tolerance + atol : float + Absolute tolerance + direction : float + Integration direction + K : tuple[[float,...],...] + Storage array for RK stages + rr_old : tuple[float,float,float] + Last state 0:3 + vv_old : tuple[float,float,float] + Last state 3:6 + t_old : float + Last time + h_previous : float + Last step length + status : float + Solver status + fr : tuple[float,float,float] + Current value of the derivative 0:3 + fv : tuple[float,float,float] + Current value of the derivative 3:6 + h_abs : float + Absolute step + + Returns + ------- + t : float + Current time. + rr : tuple[float,float,float] + Current state 0:3 + vv : tuple[float,float,float] + Current state 3:6 + t_bound : float + Boundary time + fun : Callable + Right-hand side of the system + argk : float + Standard gravitational parameter for `fun` + rtol : float + Relative tolerance + atol : float + Absolute tolerance + direction : float + Integration direction + K : tuple[[float,...],...] + Storage array for RK stages + rr_old : tuple[float,float,float] + Last state 0:3 + vv_old : tuple[float,float,float] + Last state 3:6 + t_old : float + Last time + h_previous : float + Last step length + status : float + Solver status + fr : tuple[float,float,float] + Current value of the derivative 0:3 + fv : tuple[float,float,float] + Current value of the derivative 3:6 + h_abs : float + Absolute step + """ if status != DOP853_RUNNING: From eff4f8af1aea2d2caf8513bfb609f2363233c89e Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Apr 2024 19:39:06 +0200 Subject: [PATCH 345/346] doc string --- src/hapsira/core/math/ivp/_rkdenseinterp.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/hapsira/core/math/ivp/_rkdenseinterp.py b/src/hapsira/core/math/ivp/_rkdenseinterp.py index 790300258..2709b59fd 100644 --- a/src/hapsira/core/math/ivp/_rkdenseinterp.py +++ b/src/hapsira/core/math/ivp/_rkdenseinterp.py @@ -15,6 +15,9 @@ def dop853_dense_interp_hf(t, t_old, h, rr_old, vv_old, F): Local interpolant over step made by an ODE solver. Evaluate the interpolant. + Based on + https://github.com/scipy/scipy/blob/4edfcaa3ce8a387450b6efce968572def71be089/scipy/integrate/_ivp/rk.py#L584 + Parameters ---------- t : float or array_like with shape (n_points,) From 97b4c930f512de5aaa11bc96cbe648fc2e6a55fa Mon Sep 17 00:00:00 2001 From: "Sebastian M. Ernst" Date: Mon, 29 Apr 2024 19:49:17 +0200 Subject: [PATCH 346/346] doc string --- src/hapsira/core/math/ivp/_rkdenseinterp.py | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/src/hapsira/core/math/ivp/_rkdenseinterp.py b/src/hapsira/core/math/ivp/_rkdenseinterp.py index 2709b59fd..6b675b686 100644 --- a/src/hapsira/core/math/ivp/_rkdenseinterp.py +++ b/src/hapsira/core/math/ivp/_rkdenseinterp.py @@ -20,14 +20,25 @@ def dop853_dense_interp_hf(t, t_old, h, rr_old, vv_old, F): Parameters ---------- - t : float or array_like with shape (n_points,) - Points to evaluate the solution at. + t : float + Current time. + t_old : float + Previous time. + h : float + Step to use. + rr_rold : tuple[float,float,float] + Last values 0:3. + vv_vold : tuple[float,float,float] + Last values 3:6. + F : tuple[tuple[float,...]...] + Dense output coefficients. Returns ------- - y : ndarray, shape (n,) or (n, n_points) - Computed values. Shape depends on whether `t` was a scalar or a - 1-D array. + rr : tuple[float,float,float] + Computed values 0:3. + vv : tuple[float,float,float] + Computed values 3:6. """ F00, F01, F02, F03, F04, F05, F06 = F