Merge pull request #192 from Fusion4Energy/fix-gvr-cyl-bug

Fix gvr cyl bug

Author: dodu94

src/f4enix/input/ww_gvr/cli.py

@@ -6,7 +6,6 @@
 from copy import deepcopy
 from enum import Enum
 from pathlib import Path
-from typing import Dict, Optional, Tuple
 
 from f4enix.input.ww_gvr import WW
 
@@ -62,7 +61,7 @@ def __init__(self) -> None:
 
         Running Menu() will start the CLI loop directly.
         """
-        self.weight_windows: Dict[str, WW] = {}
+        self.weight_windows: dict[str, WW] = {}
         self.extra_text: str = ""
         self.command_map = {
             Command.OPEN: self._handle_open,
@@ -171,7 +170,7 @@ def _handle_gvr(self):
         except FileNotFoundError:
             self.extra_text = " File not found..."
 
-    def _select_ww_key(self) -> Optional[str]:
+    def _select_ww_key(self) -> str | None:
         if len(self.weight_windows) == 0:
             self.extra_text = " No weight windows loaded..."
             return None
@@ -189,7 +188,7 @@ def _select_ww_key(self) -> Optional[str]:
                     return list(self.weight_windows.keys())[ww_key - 1]
             print(" Invalid weight window number...")
 
-    def _ask_gvr_parameters(self) -> Tuple[float, float]:
+    def _ask_gvr_parameters(self) -> tuple[float, float]:
         try:
             maximum_splitting = float(
                 input(" Enter maximum splitting ratio (default 5): ")

src/f4enix/input/ww_gvr/models.py

@@ -7,7 +7,6 @@
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
-from typing import Dict, List, Union
 
 import numpy as np
 from numpy.typing import NDArray
@@ -47,8 +46,8 @@ def __eq__(self, other):
             return False
 
 
-Pathlike = Union[str, Path]
-NestedList = List[List[float]]
-ValuesByEnergy = Dict[float, np.ndarray]
-ValuesByParticle = Dict[ParticleType, ValuesByEnergy]
-EnergiesByParticle = Dict[ParticleType, List[float]]
+Pathlike = str | Path
+NestedList = list[list[float]]
+ValuesByEnergy = dict[float, np.ndarray]
+ValuesByParticle = dict[ParticleType, ValuesByEnergy]
+EnergiesByParticle = dict[ParticleType, list[float]]

src/f4enix/input/ww_gvr/utils.py

@@ -2,14 +2,12 @@
 Some utilities for the ww_gvr package.
 """
 
-from typing import Tuple
-
 import numpy as np
 
 from f4enix.input.ww_gvr.models import Vectors
 
 
-def decompose_b2_vectors(b2_vectors: Vectors) -> Tuple[Vectors, Vectors]:
+def decompose_b2_vectors(b2_vectors: Vectors) -> tuple[Vectors, Vectors]:
     """
     Takes a Vectors object with b2 format and returns two Vectors objects, one with the
     coarse vectors and the other with the fine vectors.
@@ -32,7 +30,7 @@ def decompose_b2_vectors(b2_vectors: Vectors) -> Tuple[Vectors, Vectors]:
 
     Returns
     -------
-    Tuple[Vectors, Vectors]
+    tuple[Vectors, Vectors]
         Two Vectors objects, one with the coarse vectors and the other with the
         fine vectors.
     """

src/f4enix/input/ww_gvr/ww_parser.py

@@ -4,7 +4,7 @@
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, List, TextIO
+from typing import Any, TextIO
 
 import numpy as np
 from numpy import float64
@@ -25,20 +25,20 @@ class WWHeader:
     ni: int  # Number of particle types
     nr: int  # 10/16/16 = cartesian/cylindrical/spherical coord
     probid: str  # Date and time of run
-    ne: List[int]  # Number of energy windows of each particle
+    ne: list[int]  # Number of energy windows of each particle
     nfx: int  # Number of fine meshses in i
     nfy: int  # Number of fine meshses in j
     nfz: int  # Number of fine meshses in k
-    origin: List[float]  #  Bottom left corner for cart, bottom center for cyl
+    origin: list[float]  #  Bottom left corner for cart, bottom center for cyl
     ncx: int  # Number of coarse meshses in i
     ncy: int  # Number of coarse meshses in j
     ncz: int  # Number of coarse meshses in k
 
 
 @dataclass()
 class WWHeaderCyl(WWHeader):
-    director_1: List[float]  # Vector defining the director 1
-    director_2: List[float]  # Vector defining the director 2
+    director_1: list[float]  # Vector defining the director 1
+    director_2: list[float]  # Vector defining the director 2
 
 
 @dataclass()
@@ -102,7 +102,7 @@ def _parse_header(infile: TextIO) -> WWHeader:
     ncy = int(float(words[1]))  # Number of coarse meshes in j
     ncz = int(float(words[2]))  # Number of coarse meshes in k
 
-    header_args: Dict[str, Any] = {
+    header_args: dict[str, Any] = {
         "if_": if_,
         "iv": iv,
         "ni": ni,
@@ -147,7 +147,7 @@ def _read_block_2_vector(
      ...]
     """
     expected_length = 3 * number_of_coarse_intervals + 1
-    vector: List[float] = []
+    vector: list[float] = []
     while len(vector) < expected_length:
         words = infile.readline().split()
         vector.extend([float(word) for word in words])
@@ -162,7 +162,7 @@ def _read_block_3(infile: TextIO, header: WWHeader) -> tuple[NestedList, NestedL
     for particle_index in range(header.ni):
         # Read energy bins for this particle
         expected_energy_bins = header.ne[particle_index]
-        energies_current_particle: List[float] = []
+        energies_current_particle: list[float] = []
         while len(energies_current_particle) < expected_energy_bins:
             words = infile.readline().split()
             energies_current_particle.extend([float(word) for word in words])
@@ -171,7 +171,7 @@ def _read_block_3(infile: TextIO, header: WWHeader) -> tuple[NestedList, NestedL
         expected_value_bins = (
             header.nfx * header.nfy * header.nfz * expected_energy_bins
         )
-        values_current_particle: List[float] = []
+        values_current_particle: list[float] = []
         while len(values_current_particle) < expected_value_bins:
             words = infile.readline().split()
             values_current_particle.extend([float(word) for word in words])
@@ -224,11 +224,11 @@ def _read_header_from_meshtally_file(mesh: Fmesh) -> WWHeader:
     nfy = ncy = len(mesh.x2bin) - 1
     nfz = ncz = len(mesh.x3bin) - 1
     if mesh.trsf and any(mesh.trsf.origin):
-        origin = [mesh.trsf.origin[2], mesh.trsf.origin[1], mesh.trsf.origin[0]]
+        origin = [mesh.trsf.origin[0], mesh.trsf.origin[1], mesh.trsf.origin[2]]
     else:
         origin = [0.0, 0.0, 0.0]
 
-    header_args: Dict[str, Any] = {
+    header_args: dict[str, Any] = {
         "if_": if_,
         "iv": iv,
         "ni": ni,
@@ -257,6 +257,12 @@ def _read_header_from_meshtally_file(mesh: Fmesh) -> WWHeader:
                 director_2 = [0.0, 1.0, 0.0]  # Cant be the same as the axis
         else:
             director_2 = mesh.trsf.vec.tolist()
+        # The director vectors are AXS and VEC but with origin in the bottom center of
+        # the cylinder instead of 0,0,0 as provided by MESHTAL.
+        height = mesh.x2bin[-1] - mesh.x2bin[0]
+        director_1 = (np.array(director_1) * height + np.array(origin)).tolist()
+        radius = mesh.x1bin[-1] - mesh.x1bin[0]
+        director_2 = (np.array(director_2) * radius + np.array(origin)).tolist()
         header_args.update(
             {
                 "director_1": director_1,

tests/test_ww_gvr/test_ww_parser.py

@@ -189,16 +189,16 @@ def test_read_meshtally_file_cyl():
         ncx=3,
         ncy=10,
         ncz=10,
-        director_1=[1.0, 0.0, 0.0],
-        director_2=[0.0, 0.0, 1.0],
+        director_1=[0.0, -50.0, -50.0],
+        director_2=[-50.0, -50.0, 0.0],
     )
     expected_b2_vec_i = np.array(
         [0.0, 1.0, 16.67, 1.0, 1.0, 33.33, 1.0, 1.0, 50.0, 1.0]
     )
     expected_energies = [[100]]
 
     result = read_meshtally_file(
-        Path("tests") / "test_ww_gvr" / "resources" / "meshtal_cyl",  tally_id=4
+        Path("tests") / "test_ww_gvr" / "resources" / "meshtal_cyl", tally_id=4
     )
 
     assert expected_ww_header == result.header