Skip to content

Commit

Permalink
Fix parallelization in vacuum synthesis
Browse files Browse the repository at this point in the history
Change-Id: I62eebd9acaffe302a6fe9f5f38978d461dcaf553
  • Loading branch information
adrien-berchet committed Nov 25, 2020
1 parent d983444 commit e2c68dc
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 16 deletions.
6 changes: 2 additions & 4 deletions src/synthesis_workflow/tasks/synthesis.py
Original file line number Diff line number Diff line change
Expand Up @@ -338,10 +338,8 @@ def run(self):
destination=annotations_file,
)
else:
annotations_file = (
Path(PathConfig().local_synthesis_input_path)
/ self.annotations_path
)
input_task_target = yield GetSynthesisInputs()
annotations_file = input_task_target.ppath / self.annotations_path
axon_cells = None
neurondb_path = find_case_insensitive_file(self.get_neuron_db_path("dat"))

Expand Down
34 changes: 22 additions & 12 deletions src/synthesis_workflow/vacuum_synthesis.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,10 @@
"""Functions for synthesis to be used by luigi tasks."""
from functools import partial

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from joblib import cpu_count
from joblib import delayed
from joblib import Parallel
from matplotlib.backends.backend_pdf import PdfPages
Expand Down Expand Up @@ -52,6 +55,15 @@ def _grow_morphology(
return vacuum_synth_morphs_df


def _external_diametrizer(neuron, model, neurite_type, diameter_params=None):
return build_diameters.build(
neuron,
model,
[neurite_type],
diameter_params,
)


def grow_vacuum_morphologies(
mtypes,
n_cells,
Expand All @@ -75,23 +87,21 @@ def grow_vacuum_morphologies(
tmd_distributions["mtypes"][mtype]["diameter"]["method"] = diametrizer

if diametrizer == "external":

def external_diametrizer(neuron, model, neurite_type):
return build_diameters.build(
neuron,
model,
[neurite_type],
tmd_parameters[mtype][ # pylint: disable=cell-var-from-loop
"diameter_params"
],
)

external_diametrizer = partial(
_external_diametrizer,
diameter_params=tmd_parameters[mtype]["diameter_params"],
)
else:
external_diametrizer = None

gids = range(global_gid, global_gid + n_cells)
global_gid += n_cells
for row in Parallel(nb_jobs, verbose=joblib_verbose)(
for row in Parallel(
nb_jobs,
verbose=joblib_verbose,
backend="multiprocessing",
batch_size=1 + int(len(gids) / (cpu_count() if nb_jobs == -1 else nb_jobs)),
)(
delayed(_grow_morphology)(
gid,
mtype,
Expand Down

0 comments on commit e2c68dc

Please sign in to comment.