diff --git a/atlas_densities/densities/mtype_densities_from_map/create.py b/atlas_densities/densities/mtype_densities_from_map/create.py index 936005d..50f1b80 100644 --- a/atlas_densities/densities/mtype_densities_from_map/create.py +++ b/atlas_densities/densities/mtype_densities_from_map/create.py @@ -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, @@ -95,34 +96,33 @@ 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() - } + region_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] + metype_density = region_index.ravel(np.zeros(annotation.shape, dtype=float)) + for region_acronym, region_id in region_info.itertuples(index=False): + + region_indices = region_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 = region_index.ravel(molecular_type_densities[molecular_type]) + metype_density[region_indices] += density[region_indices] * coefficient + metype_density = np.reshape(metype_density, annotation.shape, order=region_index._order) if np.any(metype_density): # save density file @@ -137,7 +137,7 @@ def _create_densities_for_metype(metype: str) -> Optional[Tuple[str, str]]: returns = Parallel(n_jobs=n_jobs, return_as="generator")( delayed(_create_densities_for_metype)(metype) for metype in probability_map.columns ) - + #returns = [_create_densities_for_metype(metype) for metype in probability_map.columns[0:10]] # construct metadata output_legend: Dict[str, Dict[str, str]] = defaultdict(dict) for return_value in tqdm(returns, total=len(probability_map.columns)):