Added usrbin.py (XYZ, RZ)

usrbin.py

@@ -0,0 +1,350 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+usrbin.py — A FLUKA USRBIN ASCII reader.
+
+Author
+------
+Giuseppe Magro
+Medical Physicist, CNAO Foundation (Pavia, Italy)
+Email: [email protected]
+
+Description
+-----------
+This module provides a fully-featured reader for FLUKA USRBIN ASCII output
+files, supporting both Cartesian and R–Z binning geometries.
+
+The reader parses:
+    - Mesh type (Cartesian or R–Z)
+    - Mesh name
+    - Particle type
+    - Binning structure (min/max, bin width, number of bins)
+    - Scored data values
+    - Associated percentage errors
+
+All parsed quantities are exposed using lightweight `AttributeDict` objects,
+providing both attribute-style access (e.g. `obj.X.Min`) and dictionary-style
+access.
+
+Data reshaping follows Fortran ordering ("F"), in accordance with FLUKA.
+
+Save/load support:
+The module allows the fully parsed mesh to be exported to a compact pickle
+(.pkl) file and reloaded later without the need to re-parse the original ASCII.
+
+Example
+-------
+Reading a USRBIN file:
+
+    from usrbin import ReadFile
+
+    usr = ReadFile("dose_25.txt")
+
+    print(usr.MeshInfo)       # mesh name and type
+    print(usr.Particle)       # particle type
+    print(usr.Binning.X.Min)  # access X binning info
+
+    # Access bin centers
+    mesh_mid = usr.GetMesh("mid")
+    print(mesh_mid.X)
+
+    # Access scored data and associated uncertainties
+    data = usr.GetData()
+    errors = usr.GetErrors()
+
+Saving to pickle:
+
+    usr.Save()  # creates 'dose_25.pkl'
+
+Loading the saved pickle:
+
+    from usrbin import ReadFile
+    saved = ReadFile.Load("dose_25.pkl")
+
+    print(saved.MeshInfo.Type)
+    print(saved.Data.shape)
+
+Notes
+-----
+This reader supports only Cartesian and R–Z geometries.
+Additional geometries may be added in future extensions.
+
+License
+-------
+MIT License (or update with your preferred license).
+"""
+
+import numpy as np
+import pickle
+from typing import Tuple, Dict, Any
+
+
+# =============================================================================
+# Utility structure
+# =============================================================================
+
+class AttributeDict(dict):
+    """
+    A dictionary allowing attribute-style access.
+
+    Example:
+        obj = AttributeDict({'a': 1})
+        print(obj.a)        # 1
+        obj.b = 2
+        print(obj['b'])     # 2
+    """
+
+    __getattr__ = dict.get
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+    def __getstate__(self):
+        return self.__dict__
+
+    def __setstate__(self, d):
+        self.__dict__.update(d)
+
+
+# =============================================================================
+# USRBIN Reader
+# =============================================================================
+
+class ReadFile:
+    """
+    Read a FLUKA USRBIN ASCII file and expose mesh metadata, binning and data.
+
+    After reading, accessible fields include:
+        - MeshInfo      {'Type', 'Name'}
+        - Particle      (string)
+        - Binning       AttributeDict containing X/Y/Z or R/Z/AxisCoordinates
+        - Data          numpy array, reshaped in Fortran order ("F")
+        - Errors        numpy array, percentage errors
+    """
+
+    # -------------------------------------------------------------------------
+    def __init__(self, filename: str):
+        self.FileName = filename
+
+        with open(filename, "r") as f:
+            line = f.readline()
+
+            # Skip preliminary header lines beginning with "1"
+            while line and line.startswith("1"):
+                line = f.readline()
+
+                try:
+                    MeshTypeRaw, MeshName, Particle = line.split('"')
+                except ValueError:
+                    raise ValueError(f"Malformed mesh descriptor line:\n{line}")
+
+                MeshType = self._normalize_mesh_type(MeshTypeRaw)
+                MeshName = MeshName.strip()
+                Particle = Particle.split()[-1]
+
+                if MeshType not in ("Cartesian", "R - Z"):
+                    raise ValueError(
+                        "Only <Cartesian> and <R - Z> binning are currently supported."
+                    )
+
+                # ------------------------ Read binning header -------------------------
+                obj1, obj2, obj3 = self.ReadHeader(MeshType, f)
+
+                if MeshType == "Cartesian":
+                    X, Y, Z = obj1, obj2, obj3
+                    N = X.Nbins * Y.Nbins * Z.Nbins
+                else:
+                    R, Z, AxisCoordinates = obj1, obj2, obj3
+                    N = R.Nbins * Z.Nbins
+
+                # Skip four descriptor lines before the data
+                for _ in range(4):
+                    f.readline()
+
+                # -------------------------- Read data array ---------------------------
+                Data = self._read_float_array(f, N)
+
+                # Skip blank lines and error header
+                f.readline()
+                f.readline()
+                f.readline()
+
+                # ----------------------- Read percentage errors -----------------------
+                Errors = self._read_float_array(f, N)
+
+                # ----------------------- Store metadata ------------------------------
+                self.MeshInfo = AttributeDict({"Type": MeshType, "Name": MeshName})
+                self.Particle = Particle
+
+                if MeshType == "Cartesian":
+                    self.X, self.Y, self.Z = X, Y, Z
+                    self.Binning = AttributeDict({"X": X, "Y": Y, "Z": Z})
+
+                    self.Data = Data.reshape(
+                        (X.Nbins, Y.Nbins, Z.Nbins), order="F"
+                    )
+                    self.Errors = Errors.reshape(
+                        (X.Nbins, Y.Nbins, Z.Nbins), order="F"
+                    )
+
+                else:  # R - Z
+                    self.R, self.Z, self.AxisCoordinates = R, Z, AxisCoordinates
+                    self.Binning = AttributeDict(
+                        {"R": R, "Z": Z, "AxisCoordinates": AxisCoordinates}
+                    )
+
+                    self.Data = Data.reshape((R.Nbins, Z.Nbins), order="F")
+                    self.Errors = Errors.reshape((R.Nbins, Z.Nbins), order="F")
+
+    # -------------------------------------------------------------------------
+    @staticmethod
+    def _normalize_mesh_type(raw: str) -> str:
+        """Normalize mesh type from the raw descriptor."""
+        tokens = raw.split()
+        if "".join(tokens[:3]) == "R-Z":
+            return "R - Z"
+        return tokens[0]
+
+    # -------------------------------------------------------------------------
+    @staticmethod
+    def _read_float_array(f, N: int) -> np.ndarray:
+        """Read N floating-point values from possibly multi-line blocks."""
+        values = []
+        while len(values) < N:
+            line = f.readline()
+            if not line:
+                raise EOFError("Unexpected EOF while reading USRBIN data.")
+            values.extend(float(x) for x in line.split())
+        return np.array(values)
+
+    # -------------------------------------------------------------------------
+    def ReadHeader(self, MeshType: str, f):
+        """Read FLUKA binning header depending on geometry."""
+        if MeshType == "Cartesian":
+            X = AttributeDict(self._parse_cartesian(f))
+            Y = AttributeDict(self._parse_cartesian(f))
+            Z = AttributeDict(self._parse_cartesian(f))
+            return X, Y, Z
+
+        # R - Z
+        R = AttributeDict(self._parse_cartesian(f))
+        Z = AttributeDict(self._parse_cartesian(f))
+        X0, Y0 = self._parse_axis(f)
+        Axis = AttributeDict({"X": X0, "Y": Y0})
+        return R, Z, Axis
+
+    # -------------------------------------------------------------------------
+    def _parse_cartesian(self, f) -> Dict[str, Any]:
+        """Parse a Cartesian binning descriptor line."""
+        tokens = f.readline().split()
+        return {
+            "Min": float(tokens[3]),
+            "Max": float(tokens[5]),
+            "Nbins": int(tokens[7]),
+            "Width": float(tokens[10]),
+        }
+
+    # -------------------------------------------------------------------------
+    def _parse_axis(self, f) -> Tuple[float, float]:
+        """Parse axis coordinates line for R–Z geometry."""
+        tokens = f.readline().split()
+        try:
+            X0 = float(tokens[4].replace(",", ""))
+            Y0 = float(tokens[7])
+        except Exception:
+            raise ValueError(f"Error parsing axis coordinates line:\n{tokens}")
+        return X0, Y0
+
+    # -------------------------------------------------------------------------
+    def GetMesh(self, BinEdge: str = "mid") -> AttributeDict:
+        """
+        Return coordinate arrays for either bin centers ('mid')
+        or lower bin edges ('low').
+        """
+        if BinEdge not in ("mid", "low"):
+            raise ValueError("BinEdge must be 'mid' or 'low'.")
+
+        if self.MeshInfo.Type == "Cartesian":
+            return self._get_mesh_cartesian(BinEdge)
+        return self._get_mesh_rz(BinEdge)
+
+    # -------------------------------------------------------------------------
+    def _get_mesh_cartesian(self, edge: str) -> AttributeDict:
+        X, Y, Z = self.X, self.Y, self.Z
+
+        if edge == "mid":
+            Xg = np.linspace(X.Min + 0.5 * X.Width, X.Max - 0.5 * X.Width, X.Nbins)
+            Yg = np.linspace(Y.Min + 0.5 * Y.Width, Y.Max - 0.5 * Y.Width, Y.Nbins)
+            Zg = np.linspace(Z.Min + 0.5 * Z.Width, Z.Max - 0.5 * Z.Width, Z.Nbins)
+        else:
+            Xg = np.linspace(X.Min, X.Max, X.Nbins + 1)
+            Yg = np.linspace(Y.Min, Y.Max, Y.Nbins + 1)
+            Zg = np.linspace(Z.Min, Z.Max, Z.Nbins + 1)
+
+        return AttributeDict({"X": Xg, "Y": Yg, "Z": Zg})
+
+    # -------------------------------------------------------------------------
+    def _get_mesh_rz(self, edge: str) -> AttributeDict:
+        R, Z = self.R, self.Z
+
+        if edge == "mid":
+            Rg = np.linspace(R.Min + 0.5 * R.Width, R.Max - 0.5 * R.Width, R.Nbins)
+            Zg = np.linspace(Z.Min + 0.5 * Z.Width, Z.Max - 0.5 * Z.Width, Z.Nbins)
+        else:
+            Rg = np.linspace(R.Min, R.Max, R.Nbins + 1)
+            Zg = np.linspace(Z.Min, Z.Max, Z.Nbins + 1)
+
+        return AttributeDict({"R": Rg, "Z": Zg})
+
+    # -------------------------------------------------------------------------
+    # Simple getters
+    # -------------------------------------------------------------------------
+    def GetMeshInfo(self):
+        return self.MeshInfo
+
+    def GetParticle(self):
+        return self.Particle
+
+    def GetBinning(self):
+        return self.Binning
+
+    def GetData(self):
+        return self.Data
+
+    def GetErrors(self):
+        return self.Errors
+
+    # -------------------------------------------------------------------------
+    def Save(self, output_name: str = None):
+        """
+        Save parsed data and metadata into a pickle file.
+        """
+        if output_name is None:
+            output_name = self.FileName.replace(".txt", ".pkl")
+
+        payload = AttributeDict(
+            {
+                "MeshInfo": self.MeshInfo,
+                "Particle": self.Particle,
+                "Binning": self.Binning,
+                "Data": self.Data,
+                "Errors": self.Errors,
+            }
+        )
+
+        with open(output_name, "wb") as f:
+            pickle.dump(payload, f)
+
+    # -------------------------------------------------------------------------
+    @staticmethod
+    def Load(pkl_file: str):
+        """
+        Load a previously saved USRBIN pickle (.pkl) file.
+        """
+        with open(pkl_file, "rb") as f:
+            return pickle.load(f)
+
+    # -------------------------------------------------------------------------
+    def __repr__(self):
+        """String representation for debugging."""
+        return (f"<USRBIN Mesh: {self.MeshInfo.Name} "
+                f"({self.MeshInfo.Type}), Particle={self.Particle}>")