Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feature/parallel initializer #997

Merged
merged 4 commits into from
Apr 15, 2024
Merged

feature/parallel initializer #997

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e2dfd26
implemented batched, sneaky pool initialization
rhayes777 Apr 15, 2024
b9ddc12
fixing tests
rhayes777 Apr 15, 2024
c88347c
fixing initializer tests...
rhayes777 Apr 15, 2024
d296a74
ensure number of cores and paths are passed to initialiser
rhayes777 Apr 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 47 additions & 22 deletions autofit/non_linear/initializer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,15 @@
import os
import random
from abc import ABC, abstractmethod
from typing import Dict, Tuple, List
from typing import Dict, Tuple, List, Optional

import numpy as np

from autofit import exc
from autofit.non_linear.paths.abstract import AbstractPaths
from autofit.mapper.prior.abstract import Prior
from autofit.mapper.prior_model.abstract import AbstractPriorModel
from autofit.non_linear.parallel import SneakyPool

logger = logging.getLogger(__name__)

Expand All @@ -23,13 +25,28 @@ class AbstractInitializer(ABC):
def _generate_unit_parameter_list(self, model):
pass

@staticmethod
def figure_of_metric(args) -> Optional[float]:
fitness, parameter_list = args
try:
figure_of_merit = fitness(parameters=parameter_list)

if np.isnan(figure_of_merit) or figure_of_merit < -1e98:
return None

return figure_of_merit
except exc.FitException:
return None

def samples_from_model(
self,
total_points: int,
model: AbstractPriorModel,
fitness,
paths: AbstractPaths,
use_prior_medians: bool = False,
test_mode_samples: bool = True,
n_cores: int = 1,
):
"""
Generate the initial points of the non-linear search, by randomly drawing unit values from a uniform
Expand All @@ -55,31 +72,39 @@ def samples_from_model(
parameter_lists = []
figures_of_merit_list = []

point_index = 0

while point_index < total_points:
if not use_prior_medians:
unit_parameter_list = self._generate_unit_parameter_list(model)
else:
unit_parameter_list = [0.5] * model.prior_count

parameter_list = model.vector_from_unit_vector(
unit_vector=unit_parameter_list
)
sneaky_pool = SneakyPool(n_cores, fitness, paths)

try:
figure_of_merit = fitness(parameters=parameter_list)
while len(figures_of_merit_list) < total_points:
remaining_points = total_points - len(figures_of_merit_list)
batch_size = min(remaining_points, n_cores)
parameter_lists_ = []
unit_parameter_lists_ = []

if np.isnan(figure_of_merit) or figure_of_merit < -1e98:
raise exc.FitException
for _ in range(batch_size):
if not use_prior_medians:
unit_parameter_list = self._generate_unit_parameter_list(model)
else:
unit_parameter_list = [0.5] * model.prior_count

unit_parameter_lists.append(unit_parameter_list)
parameter_lists.append(parameter_list)
figures_of_merit_list.append(figure_of_merit)
point_index += 1
except exc.FitException:
pass
parameter_list = model.vector_from_unit_vector(
unit_vector=unit_parameter_list
)

parameter_lists_.append(parameter_list)
unit_parameter_lists_.append(unit_parameter_list)

for figure_of_merit, unit_parameter_list, parameter_list in zip(
sneaky_pool.map(
self.figure_of_metric,
[(fitness, parameter_list) for parameter_list in parameter_lists_],
),
unit_parameter_lists_,
parameter_lists_,
):
if figure_of_merit is not None:
unit_parameter_lists.append(unit_parameter_list)
parameter_lists.append(parameter_list)
figures_of_merit_list.append(figure_of_merit)

if total_points > 1 and np.allclose(
a=figures_of_merit_list[0], b=figures_of_merit_list[1:]
Expand Down
30 changes: 17 additions & 13 deletions autofit/non_linear/search/mcmc/emcee/search.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ def _fit(self, model: AbstractPriorModel, analysis):
analysis=analysis,
paths=self.paths,
fom_is_log_likelihood=False,
resample_figure_of_merit=-np.inf
resample_figure_of_merit=-np.inf,
)

pool = self.make_sneaky_pool(fitness)
Expand Down Expand Up @@ -147,6 +147,8 @@ def _fit(self, model: AbstractPriorModel, analysis):
total_points=search_internal.nwalkers,
model=model,
fitness=fitness,
paths=self.paths,
n_cores=self.number_of_cores,
)

state = np.zeros(shape=(search_internal.nwalkers, model.prior_count))
Expand Down Expand Up @@ -184,17 +186,19 @@ def _fit(self, model: AbstractPriorModel, analysis):
samples = self.samples_from(model=model, search_internal=search_internal)

if self.auto_correlation_settings.check_for_convergence:
if search_internal.iteration > self.auto_correlation_settings.check_size:
if (
search_internal.iteration
> self.auto_correlation_settings.check_size
):
if samples.converged:
iterations_remaining = 0

if iterations_remaining > 0:

self.perform_update(
model=model,
analysis=analysis,
search_internal=search_internal,
during_analysis=True
during_analysis=True,
)

return search_internal
Expand All @@ -214,7 +218,6 @@ def output_search_internal(self, search_internal):
pass

def samples_info_from(self, search_internal=None):

search_internal = search_internal or self.backend

auto_correlations = self.auto_correlations_from(search_internal=search_internal)
Expand All @@ -225,7 +228,7 @@ def samples_info_from(self, search_internal=None):
"change_threshold": auto_correlations.change_threshold,
"total_walkers": len(search_internal.get_chain()[0, :, 0]),
"total_steps": len(search_internal.get_log_prob()),
"time": self.timer.time if self.timer else None
"time": self.timer.time if self.timer else None,
}

def samples_via_internal_from(self, model, search_internal=None):
Expand All @@ -247,14 +250,14 @@ def samples_via_internal_from(self, model, search_internal=None):
search_internal = search_internal or self.backend

if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":

samples_after_burn_in = search_internal.get_chain(
discard=5, thin=5, flat=True
)
discard=5, thin=5, flat=True
)

else:

auto_correlations = self.auto_correlations_from(search_internal=search_internal)
auto_correlations = self.auto_correlations_from(
search_internal=search_internal
)

discard = int(3.0 * np.max(auto_correlations.times))
thin = int(np.max(auto_correlations.times) / 2.0)
Expand Down Expand Up @@ -292,11 +295,12 @@ def samples_via_internal_from(self, model, search_internal=None):
sample_list=sample_list,
samples_info=self.samples_info_from(search_internal=search_internal),
auto_correlation_settings=self.auto_correlation_settings,
auto_correlations=self.auto_correlations_from(search_internal=search_internal),
auto_correlations=self.auto_correlations_from(
search_internal=search_internal
),
)

def auto_correlations_from(self, search_internal=None):

search_internal = search_internal or self.backend

times = search_internal.get_autocorr_time(tol=0)
Expand Down
Loading
Loading