Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-05-27 07:24:18

0001 # This file is part of the ACTS project.
0002 #
0003 # Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 #
0005 # This Source Code Form is subject to the terms of the Mozilla Public
0006 # License, v. 2.0. If a copy of the MPL was not distributed with this
0007 # file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008 
0009 # python includes
0010 import copy
0011 import json
0012 import math
0013 import sys
0014 
0015 """Append a surface to the volume and save the new surface index"""
0016 
0017 
0018 def __append_surface(volume, surface, extent, new_idx, idx_dict):
0019     # If the 'mask' entry has not been updated, remove it now
0020     if "mask" in surface:
0021         surface["masks"] = [surface["mask"]]
0022         del surface["mask"]
0023 
0024     # Map the original surface index to the new one after merging
0025     idx_dict[surface["index_in_coll"]] = (new_idx, extent)
0026     surface["index_in_coll"] = new_idx
0027 
0028     # Add the updated surface to the volume
0029     volume["surfaces"].append(surface)
0030 
0031     # Update and return the current surface index
0032     new_idx += 1
0033     return new_idx
0034 
0035 
0036 """ Add a new value to a dictionary of lists. Append if the key exists"""
0037 
0038 
0039 def __add_to_dict(key, value, dictionary):
0040     if key in dictionary:
0041         dictionary[key].append(value)
0042     else:
0043         dictionary[key] = [value]
0044 
0045 
0046 """ Read geometry data from json and merge surfaces """
0047 
0048 
0049 def merge_surfaces(logging, in_json):
0050 
0051     # Copy identical data
0052     out_json = {}
0053     out_json["header"] = in_json["header"]
0054     logging.info("Geometry: Converted header")
0055 
0056     out_json["data"] = {}
0057     out_json["data"]["volumes"] = []
0058 
0059     # Map old to new indices per volume
0060     sf_index_per_volume = {}
0061     # Total number of surfaces after merging
0062     n_surfaces = 0
0063 
0064     # Go through each volume
0065     for volume in in_json["data"]["volumes"]:
0066 
0067         sf_index_per_volume[volume["index"]] = {}
0068         sf_index_dict = sf_index_per_volume[volume["index"]]
0069 
0070         # Save the input surfaces
0071         old_surfaces = volume["surfaces"]
0072         volume["surfaces"] = []
0073 
0074         # Sort the portals by shape for merging
0075         shape_dict = {}
0076 
0077         # Go through the surfaces
0078         sf_idx = 0
0079         for surface in old_surfaces:
0080             # Do NOT merge sensitive or passive surfaces
0081             if surface["type"] == 1 or surface["type"] == 2:
0082                 sf_idx = __append_surface(
0083                     volume, surface, [0, 0], sf_idx, sf_index_dict
0084                 )
0085                 continue
0086 
0087             # Add portals to shape dict
0088             if "mask" in surface:
0089                 shape_id = surface["mask"]["shape"]
0090                 __add_to_dict(shape_id, surface, shape_dict)
0091             elif "masks" in surface:
0092                 shape_id = surface["masks"][0]["shape"]
0093                 __add_to_dict(shape_id, surface, shape_dict)
0094 
0095         # Create new surfaces of same shape by merging masks into list
0096 
0097         # Identify the rings that can be merged to a disc, by comparing the
0098         # z-position (portals do not have rotations or x, y translation)
0099         ring_dict = {}
0100         for ring in shape_dict[6]:
0101             z_pos = ring["transform"]["translation"][2]
0102             __add_to_dict(z_pos, ring, ring_dict)
0103 
0104         for rings in ring_dict.values():
0105             # Nothing to merge, just append to output
0106             if len(rings) == 1:
0107                 boundaries = (
0108                     rings[0]["mask"]["boundaries"]
0109                     if "mask" in rings[0]
0110                     else rings[0]["masks"][0]["boundaries"]
0111                 )
0112                 extent = [boundaries[0], boundaries[1]]
0113                 sf_idx = __append_surface(
0114                     volume, rings[0], extent, sf_idx, sf_index_dict
0115                 )
0116 
0117                 continue
0118 
0119             # The merged disc
0120             disc = copy.deepcopy(rings[0])
0121 
0122             # Remove old mask entry and replace by masks list
0123             if "mask" in disc:
0124                 del disc["mask"]
0125                 disc["masks"] = []
0126             elif "masks" in disc and len(disc["masks"]) > 1:
0127                 logging.error(f"Geometry: surface already converted {disc}")
0128 
0129             for ring in rings:
0130                 disc["masks"].append(
0131                     ring["mask"] if "mask" in ring else ring["masks"][0]
0132                 )
0133 
0134             # Copy disc into volume
0135             boundaries = (
0136                 rings[0]["mask"]["boundaries"]
0137                 if "mask" in ring
0138                 else ring["masks"][0]["boundaries"]
0139             )
0140             extent = [boundaries[0], boundaries[1]]
0141             sf_idx = __append_surface(volume, disc, extent, sf_idx, sf_index_dict)
0142 
0143             # Add also the other surfaces to the dict
0144             for i in range(1, len(rings)):
0145                 boundaries = (
0146                     rings[i]["mask"]["boundaries"]
0147                     if "mask" in ring
0148                     else ring["masks"][0]["boundaries"]
0149                 )
0150                 extent = [boundaries[0], boundaries[1]]
0151                 sf_index_dict[rings[i]["index_in_coll"]] = (sf_idx - 1, extent)
0152 
0153         # Identify the rings that can be merged to a disc, by comparing the z-position (portals do not have rotations or x, y translation)
0154         cyl_dict = {}
0155         for cyl in shape_dict[4]:
0156             rad = (
0157                 cyl["mask"]["boundaries"][0]
0158                 if "mask" in cyl
0159                 else cyl["masks"][0]["boundaries"][0]
0160             )
0161             __add_to_dict(rad, cyl, cyl_dict)
0162 
0163         for sub_cyls in cyl_dict.values():
0164             # Set the cylinder boundaries correctly
0165             for sub_cyl in sub_cyls:
0166                 # Apply the translation of the surface to the mask z-boundaries
0167                 z_shift = sub_cyl["transform"]["translation"][2]
0168                 if "mask" in sub_cyl:
0169                     sub_cyl["mask"]["boundaries"][1] += z_shift
0170                     sub_cyl["mask"]["boundaries"][2] += z_shift
0171                 else:
0172                     sub_cyl["masks"][0]["boundaries"][1] += z_shift
0173                     sub_cyl["masks"][0]["boundaries"][2] += z_shift
0174 
0175                 sub_cyl["transform"]["translation"] = [0.0, 0.0, 0.0]
0176 
0177             # Nothing to merge, just append to output
0178             if len(sub_cyls) == 1:
0179                 boundaries = (
0180                     sub_cyls[0]["mask"]["boundaries"]
0181                     if "mask" in sub_cyl
0182                     else sub_cyls[0]["masks"][0]["boundaries"]
0183                 )
0184                 extent = [boundaries[1], boundaries[2]]
0185                 sf_idx = __append_surface(
0186                     volume, sub_cyls[0], extent, sf_idx, sf_index_dict
0187                 )
0188                 continue
0189 
0190             # The merged disc
0191             cyl = copy.deepcopy(sub_cyls[0])
0192 
0193             # Remove old mask entry and replace by masks list
0194             if "mask" in cyl:
0195                 del cyl["mask"]
0196                 cyl["masks"] = []
0197             elif "masks" in cyl and len(cyl["masks"]) > 1:
0198                 logging.error(f"Geometry: surface already converted {cyl}")
0199 
0200             # Remove translation of the merged cylinder
0201             cyl["transform"]["translation"] = [0.0, 0.0, 0.0]
0202 
0203             # Add the sub cylinders as masks
0204             for sub_cyl in sub_cyls:
0205                 cyl["masks"].append(
0206                     sub_cyl["mask"] if "mask" in sub_cyl else sub_cyl["masks"][0]
0207                 )
0208 
0209             # Extent of the first cylinder
0210             boundaries = (
0211                 sub_cyls[0]["mask"]["boundaries"]
0212                 if "mask" in sub_cyl
0213                 else sub_cyl["masks"][0]["boundaries"]
0214             )
0215             extent = [boundaries[1], boundaries[2]]
0216 
0217             # Copy disc into volume
0218             sf_idx = __append_surface(volume, cyl, extent, sf_idx, sf_index_dict)
0219 
0220             # Add also the other surfaces to the dict
0221             for i in range(1, len(sub_cyls)):
0222                 boundaries = (
0223                     sub_cyls[i]["mask"]["boundaries"]
0224                     if "mask" in sub_cyl
0225                     else sub_cyl["masks"][0]["boundaries"]
0226                 )
0227                 extent = [boundaries[1], boundaries[2]]
0228                 sf_index_dict[sub_cyls[i]["index_in_coll"]] = (sf_idx - 1, extent)
0229 
0230         out_json["data"]["volumes"].append(volume)
0231         n_surfaces += sf_idx
0232 
0233     logging.info(f"Geometry: Converted {len(out_json['data']['volumes'])} volumes")
0234 
0235     # Copy the volume search data structure
0236     if "volume_grid" in in_json["data"]:
0237         out_json["data"]["volume_grid"] = in_json["data"]["volume_grid"]
0238 
0239     logging.info("Geometry: Converted volume acceleration structure")
0240 
0241     # Update the metadata
0242     out_json["header"]["surface_count"] = n_surfaces
0243     out_json["header"]["volume_count"] = len(out_json["data"]["volumes"])
0244 
0245     return out_json, sf_index_per_volume
0246 
0247 
0248 """ Update the surface indices in accelerator grids after surface merging"""
0249 
0250 
0251 def update_grids(logging, in_json, sf_index_per_volume):
0252 
0253     # Copy identical data
0254     out_json = {}
0255     out_json["header"] = in_json["header"]
0256     logging.info("Grids: Converted header")
0257 
0258     out_json["data"] = {}
0259     out_json["data"]["grids"] = []
0260 
0261     # Go through each volume
0262     for grids in in_json["data"]["grids"]:
0263         # Go through all grids in a volume
0264         for grid in grids["grid_data"]:
0265             volume_idx = grid["owner_link"]
0266             sf_idx_map = sf_index_per_volume[volume_idx]
0267 
0268             # Update the surface indices in the grid bins
0269             for grid_bin in grid["bins"]:
0270                 for i, sf_idx in enumerate(grid_bin["content"]):
0271                     grid_bin["content"][i] = sf_idx_map[sf_idx][0]
0272 
0273         out_json["data"]["grids"].append(grids)
0274 
0275     # Update metadata (helps debugging)
0276     out_json["header"]["grid_count"] = len(out_json["data"]["grids"])
0277 
0278     logging.info(f"Grids: Converted {len(out_json['data']['grids'])} grids")
0279 
0280     return out_json
0281 
0282 
0283 """ Translate the grid bins to global position ranges, observing the translation of the second local coordinate t_y"""
0284 
0285 
0286 def __to_glob_grid(logging, mmap):
0287     axes = mmap["axes"]
0288     step_x = math.fabs(axes[0]["edges"][1] - axes[0]["edges"][0]) / axes[0]["bins"]
0289     step_y = math.fabs(axes[1]["edges"][1] - axes[1]["edges"][0]) / axes[1]["bins"]
0290 
0291     # Create a grid of the bin corner positions in global (phi, z) coordinates
0292     glob_grid = {}
0293     if mmap["grid_link"]["type"] == 3:
0294         for b_x in range(0, axes[0]["bins"]):
0295             x = b_x * step_x + axes[0]["edges"][0]
0296             for b_y in range(0, axes[1]["bins"]):
0297                 y = b_y * step_y + axes[1]["edges"][0]
0298                 bin_data = mmap["bins"][b_x * axes[1]["bins"] + b_y]
0299                 assert bin_data["loc_index"] == [b_x, b_y]
0300 
0301                 # Only one entry in material map bin
0302                 glob_grid[(x, y)] = bin_data["content"][0]
0303     # Create a grid of the bin corner positions in global (r, phi) coordinates
0304     else:
0305         for b_y in range(0, axes[1]["bins"]):
0306             y = b_y * step_y + axes[1]["edges"][0]
0307             for b_x in range(0, axes[0]["bins"]):
0308                 x = b_x * step_x + axes[0]["edges"][0]
0309                 bin_data = mmap["bins"][b_y * axes[0]["bins"] + b_x]
0310 
0311                 assert bin_data["loc_index"] == [b_x, b_y]
0312 
0313                 # Only one entry in material map bin
0314                 glob_grid[(x, y)] = bin_data["content"][0]
0315 
0316     return glob_grid
0317 
0318 
0319 """ Mix material of overlapping bins, as done in the material mapping in ACTS.
0320     See https://github.com/acts-project/acts/blob/4d3d89dd3c949cb9addd1bd507d42d1b54e58ad9/Core/src/Material/AverageMaterials.cpp"""
0321 
0322 
0323 def __mix_material(mat1, mat2):
0324     result = copy.deepcopy(mat1)
0325 
0326     # Calculate the weights according to the thickness
0327     total_t = mat1["thickness"] + mat2["thickness"]
0328     w1 = mat1["thickness"] / total_t
0329     w2 = mat2["thickness"] / total_t
0330 
0331     # Add the material slab thickness
0332     result["thickness"] = total_t
0333 
0334     # Material parametrization
0335     # 0: X0
0336     # 1: L0
0337     # 2: Ar
0338     # 3: Z
0339     # 4: Mass density
0340     # 5: Molar densitty (unused)
0341     # 6: @c material_state enum (solid, liquid, gaseous)
0342     params1 = mat1["material"]["params"]
0343     params2 = mat2["material"]["params"]
0344 
0345     if params1 != params2:
0346         # Average X0
0347         X0_inv = w1 / params1[0] + w2 / params1[0]
0348         result["material"]["params"][0] = 1 / X0_inv
0349 
0350         # Average L0
0351         L0_inv = w1 / params1[1] + w2 / params1[1]
0352         result["material"]["params"][1] = 1 / L0_inv
0353 
0354         molar_density1 = params1[5]
0355         molar_density2 = params1[5]
0356 
0357         molar_amount1 = molar_density1 * mat1["thickness"]
0358         molar_amount2 = molar_density2 * mat2["thickness"]
0359         molar_amount = molar_amount1 + molar_amount2
0360 
0361         molar_weight1 = molar_amount1 / molar_amount
0362         molar_weight2 = molar_amount2 / molar_amount
0363 
0364         # Average Ar
0365         result["material"]["params"][2] = (
0366             molar_weight1 * params1[2] + molar_weight2 * params2[2]
0367         )
0368 
0369         # Average Z
0370         z1 = params1[3]
0371         z2 = params2[3]
0372         z = 0.0
0373         if z1 > 0.0 and z2 > 0.0:
0374             z = math.exp(z1 * math.log(z1) + z2 * math.log(z2))
0375         else:
0376             z = w1 * z1 + w2 * z2
0377 
0378         result["material"]["params"][3] = z
0379 
0380         # Average mass density
0381         mass1 = params1[4] * mat1["thickness"]
0382         mass2 = params2[4] * mat2["thickness"]
0383         result["material"]["params"][4] = (mass1 + mass2) / total_t
0384 
0385         # Average molar density
0386         result["material"]["params"][5] = molar_amount / total_t
0387 
0388     return result
0389 
0390 
0391 """ Merge material maps of adjacent portals"""
0392 
0393 
0394 def __merge_material_maps(logging, mmap, new_mmap, extent, major_dir):
0395 
0396     new_glob_mmap = __to_glob_grid(logging, new_mmap)
0397 
0398     # Bin width of the merged grid
0399     axes = mmap["axes"]
0400     step_x = math.fabs(axes[0]["edges"][1] - axes[0]["edges"][0]) / axes[0]["bins"]
0401     step_y = math.fabs(axes[1]["edges"][1] - axes[1]["edges"][0]) / axes[1]["bins"]
0402 
0403     # Bin width of the grid that should be added
0404     new_axes = mmap["axes"]
0405     new_step_x = (
0406         math.fabs(new_axes[0]["edges"][1] - new_axes[0]["edges"][0])
0407         / new_axes[0]["bins"]
0408     )
0409     new_step_y = (
0410         math.fabs(new_axes[1]["edges"][1] - new_axes[1]["edges"][0])
0411         / new_axes[1]["bins"]
0412     )
0413 
0414     major_step = new_step_x if major_dir == 0 else new_step_y
0415 
0416     for pos, bin_data in new_glob_mmap.items():
0417         # Compensate the lower edge boundary check by also checking
0418         # the upper bin edge in the major direction (r for disc, z for cylinder)
0419         if (
0420             bin_data["thickness"] != 0
0421             and (
0422                 extent[0] <= pos[major_dir] or extent[0] < (pos[major_dir] + major_step)
0423             )
0424             and pos[major_dir] <= extent[1]
0425         ):
0426             # Bin index of the new position in the merged grid
0427             b_x = round((pos[0] - axes[0]["edges"][0]) / step_x + 1) - 1
0428             b_y = round((pos[1] - axes[1]["edges"][0]) / step_y + 1) - 1
0429 
0430             gbin = 0
0431             if mmap["grid_link"]["type"] == 3:
0432                 gbin = b_x * axes[1]["bins"] + b_y
0433             else:
0434                 gbin = b_y * axes[0]["bins"] + b_x
0435 
0436             if 0 <= gbin and gbin < len(mmap["bins"]):
0437                 orig_bin = mmap["bins"][gbin]["content"][0]
0438 
0439                 # Bin clash
0440                 if orig_bin["thickness"] != 0:
0441                     logging.warning(
0442                         f"Bin clash (sf {new_mmap['owner_link']}, bin {mmap['bins'][gbin]['loc_index']}): Averaging the material"
0443                     )
0444                     orig_bin = __mix_material(mat1=orig_bin, mat2=bin_data)
0445                 else:
0446                     orig_bin["thickness"] = bin_data["thickness"]
0447                     orig_bin["material"]["params"] = bin_data["material"]["params"]
0448 
0449 
0450 """ Update the surface owner indices of the material maps after surface merging"""
0451 
0452 
0453 def update_material(logging, in_json, sf_index_per_volume):
0454 
0455     # Copy identical data
0456     out_json = {}
0457     out_json["header"] = in_json["header"]
0458     logging.info("Material: Converted header")
0459 
0460     if out_json["header"]["common"]["tag"] != "material_maps":
0461         logging.error("Material: Only material maps conversion is available!")
0462         return in_json
0463 
0464     out_json["data"] = {}
0465     out_json["data"]["grids"] = []
0466 
0467     # Go through each volume grid collection
0468     for grids in in_json["data"]["grids"]:
0469         volume_idx = grids["volume_link"]
0470         sf_idx_map = sf_index_per_volume[volume_idx]
0471 
0472         # Make sure that each surface gets only one material map
0473         surface_grid_map = {}
0474 
0475         # Output grid data collection
0476         merged_grid_data = []
0477 
0478         # Set a new maximum extent and bin number for each merged grid
0479         for grid in grids["grid_data"]:
0480 
0481             # Get the surface index after merging
0482             old_sf_idx = grid["owner_link"]
0483             new_sf_idx = sf_idx_map[old_sf_idx][0]
0484 
0485             axes = grid["axes"]
0486 
0487             if new_sf_idx not in surface_grid_map:
0488                 step_x = (
0489                     math.fabs(axes[0]["edges"][1] - axes[0]["edges"][0])
0490                     / axes[0]["bins"]
0491                 )
0492                 step_y = (
0493                     math.fabs(axes[1]["edges"][1] - axes[1]["edges"][0])
0494                     / axes[1]["bins"]
0495                 )
0496 
0497                 # Set this as the merged grid
0498                 surface_grid_map[new_sf_idx] = [
0499                     copy.deepcopy(grid),
0500                     [step_x, step_y],
0501                     1,
0502                 ]
0503 
0504                 # Set the new surface index as owner link
0505                 surface_grid_map[new_sf_idx][0]["owner_link"] = new_sf_idx
0506             else:
0507                 # Increase the counter of grids to be merged
0508                 surface_grid_map[new_sf_idx][2] += 1
0509                 merged_grid = surface_grid_map[new_sf_idx][0]
0510                 merged_axes = merged_grid["axes"]
0511 
0512                 # Update the axis span and number of bins of the merged grid
0513                 merged_axes[0]["edges"][0] = min(
0514                     merged_axes[0]["edges"][0], axes[0]["edges"][0]
0515                 )
0516                 merged_axes[0]["edges"][1] = max(
0517                     merged_axes[0]["edges"][1], axes[0]["edges"][1]
0518                 )
0519                 merged_axes[1]["edges"][0] = min(
0520                     merged_axes[1]["edges"][0], axes[1]["edges"][0]
0521                 )
0522                 merged_axes[1]["edges"][1] = max(
0523                     merged_axes[1]["edges"][1], axes[1]["edges"][1]
0524                 )
0525 
0526         # For all merged grids, adapt the number of bins and set all bins to 0
0527         for sf_idx, grid_tuple in surface_grid_map.items():
0528 
0529             # Only one grid, no merging required
0530             if grid_tuple[2] == 1:
0531                 continue
0532 
0533             grid = grid_tuple[0]
0534             steps = grid_tuple[1]
0535             axes = grid["axes"]
0536 
0537             nbins_x = math.ceil(
0538                 math.fabs(axes[0]["edges"][1] - axes[0]["edges"][0]) / steps[0]
0539             )
0540             nbins_y = math.ceil(
0541                 math.fabs(axes[1]["edges"][1] - axes[1]["edges"][0]) / steps[1]
0542             )
0543 
0544             axes[0]["bins"] = nbins_x
0545             axes[1]["bins"] = nbins_y
0546 
0547             zero_bin = copy.deepcopy(grid["bins"][0])
0548             bin_content = zero_bin["content"][0]
0549             bin_content["surface_idx"] = sf_idx
0550             bin_content["thickness"] = 0
0551 
0552             # Reset bin collection and fill it with zero bins
0553             grid["bins"] = []
0554             if grid["grid_link"]["type"] == 3:
0555                 for b_x in range(0, axes[0]["bins"]):
0556                     for b_y in range(0, axes[1]["bins"]):
0557                         zero_bin["loc_index"] = [b_x, b_y]
0558                         grid["bins"].append(copy.deepcopy(zero_bin))
0559             else:
0560                 for b_y in range(0, axes[1]["bins"]):
0561                     for b_x in range(0, axes[0]["bins"]):
0562                         zero_bin["loc_index"] = [b_x, b_y]
0563                         grid["bins"].append(copy.deepcopy(zero_bin))
0564 
0565         # Add every grid to the merged grid
0566         for grid in grids["grid_data"]:
0567 
0568             old_sf_idx = grid["owner_link"]
0569 
0570             new_sf_idx = sf_idx_map[old_sf_idx][0]
0571             extent = sf_idx_map[old_sf_idx][1]
0572 
0573             # Only one grid, no merging required
0574             if surface_grid_map[new_sf_idx][2] == 1:
0575                 continue
0576 
0577             merged_grid = surface_grid_map[new_sf_idx][0]
0578 
0579             # Direction index along which grids have to be merged (cyl vs. disc)
0580             grid_type = merged_grid["grid_link"]["type"]
0581             assert grid["grid_link"]["type"] == grid_type
0582 
0583             major_index = 1 if grid_type == 3 else 0
0584 
0585             # Merge the grid by transforming to global coordinate and clipping
0586             # to surface extent
0587             __merge_material_maps(logging, merged_grid, grid, extent, major_index)
0588 
0589         # Append the merged grids
0590         for grid_tuple in surface_grid_map.values():
0591             merged_grid_data.append(grid_tuple[0])
0592 
0593         # Add the new grid collection to volume
0594         grids["grid_data"] = merged_grid_data
0595 
0596         assert len(grids["grid_data"]) <= len(surface_grid_map)
0597         out_json["data"]["grids"].append(grids)
0598 
0599     # Update metadata (helps debugging)
0600     out_json["header"]["grid_count"] = len(out_json["data"]["grids"])
0601 
0602     logging.info(f"Material: Converted {len(out_json['data']['grids'])} grids")
0603 
0604     return out_json