diff --git a/autofit/__init__.py b/autofit/__init__.py index dd64afc99..943824a7f 100644 --- a/autofit/__init__.py +++ b/autofit/__init__.py @@ -12,7 +12,6 @@ from . import mock as m from .non_linear.grid.grid_search import GridSearch as SearchGridSearch from .aggregator.base import AggBase -from .database.aggregator.aggregator import GridSearchAggregator from .graphical.expectation_propagation.history import EPHistory from .graphical.declarative.factor.analysis import AnalysisFactor from .graphical.declarative.factor.analysis import EPAnalysisFactor @@ -34,12 +33,10 @@ from .non_linear.samples import Sample from .non_linear.samples import load_from_table from .non_linear.samples import SamplesStored -from .database.aggregator import Aggregator from .aggregator.summary.aggregate_csv import AggregateCSV from .aggregator.summary.aggregate_csv import ValueType from .aggregator.summary.aggregate_images import AggregateImages from .aggregator.summary.aggregate_fits import AggregateFITS -from .database.aggregator import Query from autofit.aggregator.fit_interface import Fit from .aggregator.search_output import SearchOutput from .mapper import prior @@ -95,7 +92,6 @@ from .non_linear.search.mcmc.emcee.search import Emcee from .non_linear.search.mcmc.zeus.search import Zeus from .non_linear.search.nest.nautilus.search import Nautilus -from .non_linear.search.nest.nss.search import NSS from .non_linear.search.nest.dynesty.search.dynamic import DynestyDynamic from .non_linear.search.nest.dynesty.search.static import DynestyStatic from .non_linear.search.mle.drawer.search import Drawer @@ -136,7 +132,6 @@ from autofit.mapper.prior.arithmetic.compound import Log10 from . import example as ex -from . import database as db for type_ in ( @@ -201,3 +196,34 @@ def save_abc(pickler, obj): is_test_mode, test_mode_level, ) + +# Lazy attributes (PEP 562): NSS pulls blackjax -> jax, and the database +# aggregator pulls sqlalchemy + the declarative models — together over a +# second of import time that most sessions never use. +_LAZY_ATTRS = { + "NSS": ("autofit.non_linear.search.nest.nss.search", "NSS"), + "Aggregator": ("autofit.database.aggregator", "Aggregator"), + "Query": ("autofit.database.aggregator", "Query"), + "GridSearchAggregator": ( + "autofit.database.aggregator.aggregator", + "GridSearchAggregator", + ), + "db": ("autofit.database", None), +} + + +def __getattr__(name): + try: + module_name, attr = _LAZY_ATTRS[name] + except KeyError: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + import importlib + + module = importlib.import_module(module_name) + value = module if attr is None else getattr(module, attr) + globals()[name] = value + return value + + +def __dir__(): + return sorted(set(globals()) | set(_LAZY_ATTRS)) diff --git a/autofit/database/sqlalchemy_.py b/autofit/database/sqlalchemy_.py index 8270c1ef5..95c0fbf9e 100644 --- a/autofit/database/sqlalchemy_.py +++ b/autofit/database/sqlalchemy_.py @@ -4,44 +4,100 @@ Sufficient interface is implemented to permit import of SQLAlchemy based classes without any error. If any attempt is made to use those classes a meaningful warning is returned. + +``sa`` is a lazy proxy: the real ``sqlalchemy`` module is imported on first +attribute access, not when this module is imported, so sessions that never use +the database do not pay sqlalchemy's import cost. Modules that reference +``sa.`` in function signatures must use ``from __future__ import +annotations`` so the annotation does not trigger the import at definition time. """ -try: - import sqlalchemy as sa - from sqlalchemy.ext import declarative -except ImportError: - class MockSQlAlchemy: - def __getattr__(self, item): - return self - def __call__(self, *args, **kwargs): - return self +def fail(): + raise ImportError( + "Please install SQLAlchemy to use the database" + ) - def __mro_entries__(self, *args, **kwargs): - return tuple() - def declarative_base(self): - return MockBase +class MockBase: + def __init__(self, *args, **kwargs): + fail() + + +class MockSQlAlchemy: + def __getattr__(self, item): + return self + + def __call__(self, *args, **kwargs): + return self + + def __mro_entries__(self, *args, **kwargs): + return tuple() + + def declarative_base(self): + return MockBase + + def __getitem__(self, item): + fail() + + def __setitem__(self, key, value): + fail() + - def __getitem__(self, item): - fail() +MockBase.metadata = MockSQlAlchemy() + + +class _LazySQLAlchemy: + """ + Defers ``import sqlalchemy`` until an attribute is first accessed, then + delegates every attribute to the real module (or to ``MockSQlAlchemy`` + when sqlalchemy is not installed). + """ + + _target = None + + def _load(self): + if _LazySQLAlchemy._target is None: + try: + import sqlalchemy + + _LazySQLAlchemy._target = sqlalchemy + except ImportError: + _LazySQLAlchemy._target = MockSQlAlchemy() + return _LazySQLAlchemy._target + + def __getattr__(self, item): + return getattr(self._load(), item) + + def __call__(self, *args, **kwargs): + return self._load()(*args, **kwargs) + + def __mro_entries__(self, *args, **kwargs): + target = self._load() + if isinstance(target, MockSQlAlchemy): + return tuple() + raise TypeError(f"{target!r} cannot be used as a base class") - def __setitem__(self, key, value): - fail() + def __getitem__(self, item): + return self._load()[item] + def __setitem__(self, key, value): + self._load()[key] = value - sa = MockSQlAlchemy() - declarative = sa +class _LazyDeclarative(_LazySQLAlchemy): + def _load(self): + if _LazyDeclarative._target is None: + try: + from sqlalchemy.ext import declarative - def fail(): - raise ImportError( - "Please install SQLAlchemy to use the database" - ) + _LazyDeclarative._target = declarative + except ImportError: + _LazyDeclarative._target = MockSQlAlchemy() + return _LazyDeclarative._target - class MockBase: - def __init__(self, *args, **kwargs): - fail() +_LazyDeclarative._target = None - metadata = sa +sa = _LazySQLAlchemy() +declarative = _LazyDeclarative() diff --git a/autofit/non_linear/fitness.py b/autofit/non_linear/fitness.py index e2367bef0..878b83d1f 100644 --- a/autofit/non_linear/fitness.py +++ b/autofit/non_linear/fitness.py @@ -1,6 +1,5 @@ import logging import numpy as np -from IPython.display import clear_output import os import time @@ -440,6 +439,8 @@ def manage_quick_update(self, parameters, log_likelihood): if self.quick_update_count >= self.iterations_per_quick_update: + from IPython.display import clear_output + clear_output(wait=True) start_time = time.time() diff --git a/autofit/non_linear/grid/grid_search/result_builder.py b/autofit/non_linear/grid/grid_search/result_builder.py index eb14e4d48..0c7840ab2 100644 --- a/autofit/non_linear/grid/grid_search/result_builder.py +++ b/autofit/non_linear/grid/grid_search/result_builder.py @@ -1,9 +1,13 @@ -from typing import List, Union +from __future__ import annotations + +from typing import List, Union, TYPE_CHECKING from autofit.non_linear.paths.abstract import AbstractPaths from autofit.non_linear.samples import Samples -from autofit.database import Prior from autofit.non_linear.result import Result, Placeholder + +if TYPE_CHECKING: + from autofit.database import Prior from .job import JobResult from .result import GridSearchResult diff --git a/autofit/non_linear/paths/database.py b/autofit/non_linear/paths/database.py index 451e7e90b..6f7d03092 100644 --- a/autofit/non_linear/paths/database.py +++ b/autofit/non_linear/paths/database.py @@ -1,16 +1,18 @@ +from __future__ import annotations + import shutil -from typing import Optional, Union +from typing import Optional, Union, TYPE_CHECKING from autonerves.output import conditional_output, should_output -from autofit.database.sqlalchemy_ import sa from .abstract import AbstractPaths import numpy as np -from autofit.database.model import Fit from autonerves.dictable import to_dict, from_dict -from autofit.database.aggregator.info import Info from autofit.non_linear.samples.summary import SamplesSummary +if TYPE_CHECKING: + from autofit.database.model import Fit + class DatabasePaths(AbstractPaths): def __init__( @@ -108,6 +110,8 @@ def zip_remove(self): """ Remove files from both the symlinked folder and the output directory """ + from autofit.database.aggregator.info import Info + self.session.commit() Info(self.session).write() @@ -234,6 +238,9 @@ def remove_search_internal(self): @property def fit(self) -> Fit: + from autofit.database.model import Fit + from autofit.database.sqlalchemy_ import sa + if self._fit is None: try: self._fit = ( @@ -341,5 +348,7 @@ def save_all(self, info, *_, **kwargs): self.save_json("search", to_dict(self.search)) self.save_json("model", to_dict(self.model)) + from autofit.database.aggregator.info import Info + self.session.commit() Info(self.session).write() diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py index 92569c3c2..31fb4f4d9 100644 --- a/autofit/non_linear/search/abstract_search.py +++ b/autofit/non_linear/search/abstract_search.py @@ -16,6 +16,7 @@ import psutil if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa from autofit.non_linear.result import Result from autonerves import conf @@ -23,7 +24,6 @@ from autonerves.output import should_output from autofit import exc -from autofit.database.sqlalchemy_ import sa from autofit.graphical import ( MeanField, AnalysisFactor, diff --git a/autofit/non_linear/search/mcmc/abstract_mcmc.py b/autofit/non_linear/search/mcmc/abstract_mcmc.py index 18b788a64..77ff4489b 100644 --- a/autofit/non_linear/search/mcmc/abstract_mcmc.py +++ b/autofit/non_linear/search/mcmc/abstract_mcmc.py @@ -1,12 +1,16 @@ -from typing import Optional +from __future__ import annotations + +from typing import Optional, TYPE_CHECKING from autonerves import conf -from autofit.database.sqlalchemy_ import sa from autofit.non_linear.search.abstract_search import NonLinearSearch from autofit.non_linear.initializer import Initializer, InitializerBall from autofit.non_linear.search.mcmc.auto_correlations import AutoCorrelationsSettings from autofit.non_linear.plot import corner_cornerpy +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + class AbstractMCMC(NonLinearSearch): def __init__( diff --git a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py index 3ea0a06bc..997507cd7 100644 --- a/autofit/non_linear/search/mcmc/blackjax/nuts/search.py +++ b/autofit/non_linear/search/mcmc/blackjax/nuts/search.py @@ -1,14 +1,15 @@ +from __future__ import annotations + import logging import os import pickle from pathlib import Path -from typing import Optional +from typing import Optional, TYPE_CHECKING import numpy as np from autonerves import conf -from autofit.database.sqlalchemy_ import sa from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.fitness import Fitness from autofit.non_linear.initializer import Initializer @@ -19,6 +20,9 @@ from autofit.non_linear.samples.mcmc import SamplesMCMC from autofit.non_linear.samples.sample import Sample +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + logger = logging.getLogger(__name__) diff --git a/autofit/non_linear/search/mcmc/emcee/search.py b/autofit/non_linear/search/mcmc/emcee/search.py index 22aa7c32c..bd56f4714 100644 --- a/autofit/non_linear/search/mcmc/emcee/search.py +++ b/autofit/non_linear/search/mcmc/emcee/search.py @@ -1,13 +1,14 @@ +from __future__ import annotations + import logging import os from pathlib import Path -from typing import Dict, Optional +from typing import Dict, Optional, TYPE_CHECKING import numpy as np from autonerves import conf -from autofit.database.sqlalchemy_ import sa from autofit.mapper.model_mapper import ModelMapper from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.fitness import Fitness @@ -19,6 +20,9 @@ from autofit.non_linear.samples.sample import Sample from autofit.non_linear.samples.mcmc import SamplesMCMC +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + logger = logging.getLogger(__name__) diff --git a/autofit/non_linear/search/mcmc/zeus/search.py b/autofit/non_linear/search/mcmc/zeus/search.py index 7de5d748c..62afce9d8 100644 --- a/autofit/non_linear/search/mcmc/zeus/search.py +++ b/autofit/non_linear/search/mcmc/zeus/search.py @@ -1,10 +1,11 @@ +from __future__ import annotations + import logging -from typing import Dict, Optional +from typing import Dict, Optional, TYPE_CHECKING import numpy as np import os -from autofit.database.sqlalchemy_ import sa from autofit.mapper.model_mapper import ModelMapper from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.fitness import Fitness @@ -16,6 +17,9 @@ from autofit.non_linear.test_mode import is_test_mode from autofit.non_linear.samples.mcmc import SamplesMCMC +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + logger = logging.getLogger(__name__) diff --git a/autofit/non_linear/search/mle/bfgs/search.py b/autofit/non_linear/search/mle/bfgs/search.py index 4bfe4c365..e0471b079 100644 --- a/autofit/non_linear/search/mle/bfgs/search.py +++ b/autofit/non_linear/search/mle/bfgs/search.py @@ -1,7 +1,8 @@ -from typing import Optional +from __future__ import annotations + +from typing import Optional, TYPE_CHECKING from autofit import exc -from autofit.database.sqlalchemy_ import sa from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.search.mle.abstract_mle import AbstractMLE @@ -14,6 +15,9 @@ import numpy as np +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + class AbstractBFGS(AbstractMLE): diff --git a/autofit/non_linear/search/mle/drawer/search.py b/autofit/non_linear/search/mle/drawer/search.py index 570e73a80..f4c9550ee 100644 --- a/autofit/non_linear/search/mle/drawer/search.py +++ b/autofit/non_linear/search/mle/drawer/search.py @@ -1,7 +1,8 @@ +from __future__ import annotations + import numpy as np -from typing import Optional +from typing import Optional, TYPE_CHECKING -from autofit.database.sqlalchemy_ import sa from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.fitness import Fitness @@ -9,6 +10,9 @@ from autofit.non_linear.initializer import AbstractInitializer from autofit.non_linear.samples import Samples, Sample +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + class Drawer(AbstractMLE): __identifier_fields__ = ("total_draws",) diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index 101727d45..3c8db9127 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -1,10 +1,11 @@ +from __future__ import annotations + import inspect import pickle -from typing import Optional +from typing import Optional, TYPE_CHECKING import numpy as np -from autofit.database.sqlalchemy_ import sa from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.search.mle.abstract_mle import AbstractMLE @@ -19,6 +20,9 @@ MultiStartGradientConvergence, ) +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + class AbstractMultiStartGradient(AbstractMLE): diff --git a/autofit/non_linear/search/nest/abstract_nest.py b/autofit/non_linear/search/nest/abstract_nest.py index 127c91ce2..d86050a51 100644 --- a/autofit/non_linear/search/nest/abstract_nest.py +++ b/autofit/non_linear/search/nest/abstract_nest.py @@ -1,9 +1,10 @@ +from __future__ import annotations + from abc import ABC -from typing import Optional +from typing import Optional, TYPE_CHECKING import warnings from autonerves import conf -from autofit.database.sqlalchemy_ import sa from autofit.non_linear.search.abstract_search import NonLinearSearch from autofit.non_linear.initializer import ( InitializerPrior, @@ -12,6 +13,9 @@ ) from autofit.non_linear.plot import corner_anesthetic +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + class AbstractNest(NonLinearSearch, ABC): def __init__( diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py index 91e6301f9..0263aa6e7 100644 --- a/autofit/non_linear/search/nest/dynesty/search/abstract.py +++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py @@ -1,14 +1,15 @@ +from __future__ import annotations + import logging import os from abc import ABC from pathlib import Path -from typing import Dict, Optional, Tuple, Union +from typing import Dict, Optional, Tuple, Union, TYPE_CHECKING import numpy as np import warnings from autofit import exc -from autofit.database.sqlalchemy_ import sa from autofit.non_linear.fitness import Fitness from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.paths.null import NullPaths @@ -17,6 +18,9 @@ from autofit.non_linear.samples.nest import SamplesNest from autofit.non_linear.test_mode import is_test_mode +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + logger = logging.getLogger(__name__) diff --git a/autofit/non_linear/search/nest/dynesty/search/static.py b/autofit/non_linear/search/nest/dynesty/search/static.py index 09210f5c6..b2bb2cf44 100644 --- a/autofit/non_linear/search/nest/dynesty/search/static.py +++ b/autofit/non_linear/search/nest/dynesty/search/static.py @@ -1,15 +1,17 @@ from __future__ import annotations from pathlib import Path -from typing import Dict, Optional, Union +from typing import Dict, Optional, Union, TYPE_CHECKING -from autofit.database.sqlalchemy_ import sa from autofit.mapper.prior_model.abstract import AbstractPriorModel from .abstract import AbstractDynesty, prior_transform +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + class DynestyStatic(AbstractDynesty): __identifier_fields__ = ( "nlive", diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py index 92e905264..3ae499e8f 100644 --- a/autofit/non_linear/search/nest/nautilus/search.py +++ b/autofit/non_linear/search/nest/nautilus/search.py @@ -1,12 +1,13 @@ +from __future__ import annotations + import numpy as np import logging import os import sys from contextlib import nullcontext from pathlib import Path -from typing import Dict, Optional, Tuple +from typing import Dict, Optional, Tuple, TYPE_CHECKING -from autofit.database.sqlalchemy_ import sa from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.mapper.prior.vectorized import PriorVectorized @@ -18,6 +19,9 @@ from autofit.non_linear.samples.nest import SamplesNest from autofit.non_linear.test_mode import is_test_mode +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + logger = logging.getLogger(__name__) diff --git a/autofit/non_linear/search/nest/nss/search.py b/autofit/non_linear/search/nest/nss/search.py index 1bdb889d5..505d03194 100644 --- a/autofit/non_linear/search/nest/nss/search.py +++ b/autofit/non_linear/search/nest/nss/search.py @@ -1,12 +1,13 @@ +from __future__ import annotations + import logging import os import pickle from pathlib import Path -from typing import Optional +from typing import Optional, TYPE_CHECKING import numpy as np -from autofit.database.sqlalchemy_ import sa from autofit.mapper.prior_model.abstract import AbstractPriorModel from autofit.non_linear.fitness import Fitness from autofit.non_linear.paths.null import NullPaths @@ -15,6 +16,9 @@ from autofit.non_linear.samples.sample import Sample from autofit.non_linear.test_mode import is_test_mode +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa + try: import blackjax as _blackjax diff --git a/autofit/non_linear/settings.py b/autofit/non_linear/settings.py index 5389ab8b4..f608ecea5 100644 --- a/autofit/non_linear/settings.py +++ b/autofit/non_linear/settings.py @@ -1,6 +1,10 @@ -from typing import Optional, List +from __future__ import annotations + +from typing import Optional, List, TYPE_CHECKING + +if TYPE_CHECKING: + from autofit.database.sqlalchemy_ import sa -from autofit.database.sqlalchemy_ import sa class SettingsSearch: