Skip to content

Commit

Permalink
Speedup mtype density creation from probability maps using voxcell's …
Browse files Browse the repository at this point in the history
…ValueToIndexVoxels (#79)
  • Loading branch information
eleftherioszisis authored Apr 15, 2024
1 parent fba1e3f commit 924b6c6
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 13 deletions.
28 changes: 16 additions & 12 deletions atlas_densities/densities/mtype_densities_from_map/create.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from atlas_commons.typing import FloatArray
from joblib import Parallel, delayed
from tqdm import tqdm
from voxcell.voxel_data import ValueToIndexVoxels

from atlas_densities.densities.mtype_densities_from_map.utils import (
_check_probability_map_consistency,
Expand Down Expand Up @@ -95,34 +96,37 @@ def create_from_probability_map( # pylint: disable=too-many-arguments
.set_index("acronym")
.loc[probability_map.index.get_level_values("region")]
.reset_index()[["region", "id"]]
.drop_duplicates(subset="region")
.reset_index(drop=True)
)
region_acronyms = set(region_info.region)

region_masks = {
region_acronym: annotation.raw == region_id
for _, region_acronym, region_id in region_info.itertuples()
}
annotation_index = ValueToIndexVoxels(annotation.raw)

# ensure output directory exists
Path(output_dirpath).mkdir(exist_ok=True, parents=True)

def _create_densities_for_metype(metype: str) -> Optional[Tuple[str, str]]:
coefficients: Dict[str, Dict[str, Any]] = {}
for region_acronym in region_acronyms:
for region_acronym in region_info.region:
coefficients[region_acronym] = {
molecular_type: probability_map.at[(region_acronym, molecular_type), metype]
for molecular_type in list(molecular_type_densities.keys())
if (region_acronym, molecular_type) in probability_map.index
}

metype_density = np.zeros(annotation.shape, dtype=float)
for region_acronym in region_acronyms:
region_mask = region_masks[region_acronym]
# perform the manipulation in the 1d flat array
metype_density = np.zeros(np.prod(annotation.shape), dtype=float)

for region_acronym, region_id in region_info.itertuples(index=False):

region_indices = annotation_index.value_to_1d_indices(region_id)
for molecular_type, coefficient in coefficients[region_acronym].items():
if coefficient <= 0.0:
continue
density = molecular_type_densities[molecular_type]
metype_density[region_mask] += density[region_mask] * coefficient
density = annotation_index.ravel(molecular_type_densities[molecular_type])
metype_density[region_indices] += density[region_indices] * coefficient

# reshape the 1d metype_density array back to the annotation's shape
metype_density = annotation_index.unravel(metype_density)

if np.any(metype_density):
# save density file
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
# from the HiGHS library. We use the "highs" method in the densities module.
"scipy>=1.6.0",
"tqdm>=4.44.1",
"voxcell>=3.1.7",
"voxcell>=3.1.8", # ValueToIndexVoxels ravel/unravel
],
extras_require={
"tests": [
Expand Down

0 comments on commit 924b6c6

Please sign in to comment.