From adaa091d3bdbb35467de9b7e0ba8c87bcc676169 Mon Sep 17 00:00:00 2001 From: Christopher Albert Date: Mon, 5 Jan 2026 23:18:07 +0100 Subject: [PATCH 1/2] feat(engineering): add coil-clearance port constraint --- python/pysimple/coil_clearance.py | 392 ++++++++++++++++++++++++++++++ python/pysimple/engineering.py | 105 ++++++++ test/python/CMakeLists.txt | 2 +- test/python/test_engineering.py | 98 ++++++++ 4 files changed, 596 insertions(+), 1 deletion(-) create mode 100644 python/pysimple/coil_clearance.py diff --git a/python/pysimple/coil_clearance.py b/python/pysimple/coil_clearance.py new file mode 100644 index 00000000..bd1fa8f8 --- /dev/null +++ b/python/pysimple/coil_clearance.py @@ -0,0 +1,392 @@ +""" +Coil-aware geometric feasibility for wall ports. + +This module provides: +- Parsing of SIMPLE/libneo coil point files +- Interpolation of the wall surface from a chartmap NetCDF file +- Minimum-distance queries from wall points to coil polylines + +The intent is to enable port placement constraints that respect physical coil +clearances while optimizing against heat loads in (theta, zeta) coordinates. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +try: + import netCDF4 as nc + HAS_NETCDF4 = True +except ImportError: + HAS_NETCDF4 = False + + +def load_coil_points(coil_file: str | Path) -> tuple[np.ndarray, np.ndarray]: + """ + Load coil points from a SIMPLE/libneo-style coil file. + + Format: + First line: integer N (number of points) + Next N lines: x y z I + + Coordinates are expected to be in meters. + + Returns: + points: (N, 3) array of xyz points [m] + currents: (N,) array of currents + """ + path = Path(coil_file) + with path.open("r", encoding="utf-8") as f: + header = f.readline().strip() + if not header: + raise ValueError(f"Empty coil file: {path}") + try: + n_points = int(header.split()[0]) + except ValueError as exc: + raise ValueError(f"Invalid coil file header (expected N): {path}") from exc + + data = np.loadtxt(f, dtype=float, ndmin=2) + if data.shape[0] != n_points or data.shape[1] < 4: + raise ValueError( + f"Coil file {path} expected {n_points} rows of x y z I, got {data.shape}" + ) + + points = data[:, 0:3].astype(float, copy=False) + currents = data[:, 3].astype(float, copy=False) + return points, currents + + +def split_coil_polylines( + points: np.ndarray, + currents: np.ndarray, + *, + use_zero_current_separators: bool = True, + jump_factor: float = 25.0, +) -> list[np.ndarray]: + """ + Split concatenated coil point lists into per-coil polylines. + + Primary split rule: + current == 0 marks a coil break (row is treated as a separator). + + Fallback split rule (for files without separators): + Split when consecutive point spacing is a large jump compared to the + typical segment length (jump_factor * median_length). + """ + if points.ndim != 2 or points.shape[1] != 3: + raise ValueError("points must have shape (N, 3)") + if currents.ndim != 1 or currents.shape[0] != points.shape[0]: + raise ValueError("currents must have shape (N,)") + + polylines: list[np.ndarray] = [] + + if use_zero_current_separators and np.any(currents == 0.0): + current_poly: list[np.ndarray] = [] + for p, I in zip(points, currents, strict=True): + if I == 0.0: + if len(current_poly) >= 2: + polylines.append(np.asarray(current_poly, dtype=float)) + current_poly = [] + continue + current_poly.append(p) + if len(current_poly) >= 2: + polylines.append(np.asarray(current_poly, dtype=float)) + return polylines + + if points.shape[0] < 2: + return [] + + diffs = np.linalg.norm(points[1:] - points[:-1], axis=1) + if not np.all(np.isfinite(diffs)): + raise ValueError("Non-finite coil point coordinates encountered") + + sorted_diffs = np.sort(diffs) + cutoff = max(1, int(0.9 * sorted_diffs.size)) + typical = float(np.median(sorted_diffs[:cutoff])) + if typical <= 0.0: + raise ValueError("Degenerate coil point list (zero spacing)") + + threshold = jump_factor * typical + break_after = diffs > threshold + + start = 0 + for i, do_break in enumerate(break_after, start=0): + if do_break: + if i + 1 - start >= 2: + polylines.append(points[start:i + 1].copy()) + start = i + 1 + if points.shape[0] - start >= 2: + polylines.append(points[start:].copy()) + return polylines + + +@dataclass(frozen=True) +class CoilSegments: + """A collection of polyline segments representing coils.""" + + p0: np.ndarray # (Nseg, 3) + p1: np.ndarray # (Nseg, 3) + + @classmethod + def from_file( + cls, + coil_file: str | Path, + *, + use_zero_current_separators: bool = True, + jump_factor: float = 25.0, + ) -> "CoilSegments": + points, currents = load_coil_points(coil_file) + polylines = split_coil_polylines( + points, + currents, + use_zero_current_separators=use_zero_current_separators, + jump_factor=jump_factor, + ) + + if not polylines: + raise ValueError(f"No coil polylines found in {Path(coil_file)}") + + seg_p0: list[np.ndarray] = [] + seg_p1: list[np.ndarray] = [] + for poly in polylines: + seg_p0.append(poly[:-1]) + seg_p1.append(poly[1:]) + + p0 = np.vstack(seg_p0).astype(float, copy=False) + p1 = np.vstack(seg_p1).astype(float, copy=False) + return cls(p0=p0, p1=p1) + + +def _min_distance_point_to_segments(point: np.ndarray, p0: np.ndarray, p1: np.ndarray) -> float: + v = p1 - p0 + vv = np.einsum("ij,ij->i", v, v) + if np.any(vv == 0.0): + d2 = np.einsum("ij,ij->i", p0 - point, p0 - point) + return float(np.sqrt(np.min(d2))) + + w = point - p0 + t = np.einsum("ij,ij->i", w, v) / vv + t = np.clip(t, 0.0, 1.0) + proj = p0 + (t[:, None] * v) + d2 = np.einsum("ij,ij->i", proj - point, proj - point) + return float(np.sqrt(np.min(d2))) + + +def min_distance_points_to_segments( + points: np.ndarray, + p0: np.ndarray, + p1: np.ndarray, + *, + block_size: int = 128, +) -> np.ndarray: + """ + Compute minimum distances from many points to many segments. + + Args: + points: (M, 3) query points + p0, p1: (Nseg, 3) segment endpoints + block_size: segment blocking for memory control + + Returns: + distances: (M,) minimum distance for each query point [m] + """ + if points.ndim != 2 or points.shape[1] != 3: + raise ValueError("points must have shape (M, 3)") + if p0.shape != p1.shape or p0.ndim != 2 or p0.shape[1] != 3: + raise ValueError("p0/p1 must have shape (Nseg, 3)") + if block_size <= 0: + raise ValueError("block_size must be positive") + + m = points.shape[0] + nseg = p0.shape[0] + if m == 0: + return np.zeros((0,), dtype=float) + + if m <= 32: + return np.asarray( + [_min_distance_point_to_segments(points[i], p0, p1) for i in range(m)], + dtype=float, + ) + + min_d2 = np.full((m,), np.inf, dtype=float) + for start in range(0, nseg, block_size): + end = min(start + block_size, nseg) + p0b = p0[start:end] + p1b = p1[start:end] + + v = p1b - p0b + vv = np.einsum("ij,ij->i", v, v) + vv = np.where(vv == 0.0, 1.0, vv) + + w = points[:, None, :] - p0b[None, :, :] + t = np.einsum("mbi,bi->mb", w, v) / vv[None, :] + t = np.clip(t, 0.0, 1.0) + proj = p0b[None, :, :] + t[:, :, None] * v[None, :, :] + d2 = np.einsum("mbi,mbi->mb", points[:, None, :] - proj, points[:, None, :] - proj) + min_d2 = np.minimum(min_d2, np.min(d2, axis=1)) + + return np.sqrt(min_d2) + + +def _periodic_bilinear_indices( + grid: np.ndarray, values: np.ndarray, period: float +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + n = grid.size + if n < 2: + raise ValueError("Interpolation grid must have at least 2 points") + + v = np.mod(values, period) + i1 = np.searchsorted(grid, v, side="right") + i0 = i1 - 1 + i0 = np.mod(i0, n) + i1 = np.mod(i1, n) + + g0 = grid[i0] + g1 = grid[i1] + + wrapped = i1 <= i0 + denom = np.where(wrapped, (g1 + period) - g0, g1 - g0) + denom = np.where(denom == 0.0, 1.0, denom) + w = (v - g0) / denom + w = np.clip(w, 0.0, 1.0) + return i0, i1, w + + +@dataclass(frozen=True) +class WallSurface: + """ + Wall surface interpolation from chartmap. + + The chartmap file stores (x, y, z) on a (rho, theta, zeta) grid for one + field period. This class interpolates on the rho=1 surface and maps + arbitrary zeta in [0, 2*pi) to the correct field period by rotation. + """ + + theta: np.ndarray # (ntheta,) in [0, 2*pi) + zeta: np.ndarray # (nzeta,) in [0, 2*pi/nfp) + nfp: int + x_wall: np.ndarray # (nzeta, ntheta) [m] + y_wall: np.ndarray # (nzeta, ntheta) [m] + z_wall: np.ndarray # (nzeta, ntheta) [m] + + @classmethod + def from_chartmap(cls, chartmap_file: str | Path) -> "WallSurface": + if not HAS_NETCDF4: + raise ImportError("netCDF4 required: pip install netCDF4") + + path = Path(chartmap_file) + with nc.Dataset(path, "r") as ds: + theta = np.asarray(ds.variables["theta"][:], dtype=float) + zeta = np.asarray(ds.variables["zeta"][:], dtype=float) + nfp = int(ds.variables["num_field_periods"][...]) + + x_wall = np.asarray(ds.variables["x"][:, :, -1], dtype=float) * 0.01 + y_wall = np.asarray(ds.variables["y"][:, :, -1], dtype=float) * 0.01 + z_wall = np.asarray(ds.variables["z"][:, :, -1], dtype=float) * 0.01 + + return cls( + theta=theta, + zeta=zeta, + nfp=nfp, + x_wall=x_wall, + y_wall=y_wall, + z_wall=z_wall, + ) + + @property + def zeta_period(self) -> float: + return 2.0 * np.pi / float(self.nfp) + + def xyz(self, theta_hm: np.ndarray, zeta_full: np.ndarray) -> np.ndarray: + theta_hm = np.asarray(theta_hm, dtype=float) + zeta_full = np.asarray(zeta_full, dtype=float) + + theta_cm = np.mod(theta_hm, 2.0 * np.pi) + period = self.zeta_period + k = np.floor(zeta_full / period).astype(int) + zeta_mod = zeta_full - k * period + + th0, th1, wth = _periodic_bilinear_indices(self.theta, theta_cm, 2.0 * np.pi) + ze0, ze1, wze = _periodic_bilinear_indices(self.zeta, zeta_mod, period) + + th0, th1, wth, ze0, ze1, wze = np.broadcast_arrays(th0, th1, wth, ze0, ze1, wze) + + def interp_field(field: np.ndarray) -> np.ndarray: + f00 = field[ze0, th0] + f01 = field[ze0, th1] + f10 = field[ze1, th0] + f11 = field[ze1, th1] + f0 = (1.0 - wth) * f00 + wth * f01 + f1 = (1.0 - wth) * f10 + wth * f11 + return (1.0 - wze) * f0 + wze * f1 + + x_base = interp_field(self.x_wall) + y_base = interp_field(self.y_wall) + z_val = interp_field(self.z_wall) + + angle = k * period + ca = np.cos(angle) + sa = np.sin(angle) + x = ca * x_base - sa * y_base + y = sa * x_base + ca * y_base + + return np.stack([x, y, z_val], axis=-1) + + +@dataclass(frozen=True) +class CoilClearanceConstraint: + wall: WallSurface + coils: CoilSegments + min_clearance_m: float + + def port_violates( + self, + theta_center: float, + zeta_center: float, + theta_width: float, + zeta_width: float, + ) -> bool: + half_th = theta_width / 2.0 + half_ze = zeta_width / 2.0 + theta_pts = np.asarray( + [ + theta_center, + theta_center - half_th, + theta_center + half_th, + theta_center - half_th, + theta_center + half_th, + ], + dtype=float, + ) + zeta_pts = np.asarray( + [ + zeta_center, + zeta_center - half_ze, + zeta_center - half_ze, + zeta_center + half_ze, + zeta_center + half_ze, + ], + dtype=float, + ) + + xyz = self.wall.xyz(theta_pts, zeta_pts).reshape((-1, 3)) + d = min_distance_points_to_segments(xyz, self.coils.p0, self.coils.p1) + return bool(np.any(d < self.min_clearance_m)) + + +def clearance_map_on_heatmap_grid( + wall: WallSurface, + coils: CoilSegments, + theta_centers: np.ndarray, + zeta_centers: np.ndarray, +) -> np.ndarray: + th = np.asarray(theta_centers, dtype=float) + ze = np.asarray(zeta_centers, dtype=float) + th_grid, ze_grid = np.meshgrid(th, ze, indexing="ij") + xyz = wall.xyz(th_grid, ze_grid).reshape((-1, 3)) + d = min_distance_points_to_segments(xyz, coils.p0, coils.p1) + return d.reshape((th.size, ze.size)) + diff --git a/python/pysimple/engineering.py b/python/pysimple/engineering.py index bb645457..7939c5ee 100644 --- a/python/pysimple/engineering.py +++ b/python/pysimple/engineering.py @@ -77,6 +77,13 @@ import numpy as np +from pysimple.coil_clearance import ( + CoilClearanceConstraint, + CoilSegments, + WallSurface, + clearance_map_on_heatmap_grid, +) + try: import netCDF4 as nc HAS_NETCDF4 = True @@ -736,6 +743,7 @@ def __init__(self, heat_map: WallHeatMap): self._ports: list[dict] = [] self._exclusion_zones: list[tuple[float, float, float, float]] = [] self._max_flux_constraint: Optional[float] = None + self._coil_clearance: Optional[CoilClearanceConstraint] = None def add_port( self, @@ -780,6 +788,37 @@ def set_max_flux_constraint(self, max_flux: float) -> "PortOptimizer": self._max_flux_constraint = max_flux return self + def set_coil_clearance_constraint( + self, + coil_file: str | Path, + *, + chartmap_file: str | Path, + min_clearance_m: float, + use_zero_current_separators: bool = True, + jump_factor: float = 25.0, + ) -> "PortOptimizer": + """ + Enforce a minimum distance from ports to coils (geometric feasibility). + + This constraint requires a chartmap file to map (theta, zeta) port + locations to physical-space coordinates on the wall. + """ + if min_clearance_m <= 0.0: + raise ValueError("min_clearance_m must be positive") + + wall = WallSurface.from_chartmap(chartmap_file) + coils = CoilSegments.from_file( + coil_file, + use_zero_current_separators=use_zero_current_separators, + jump_factor=jump_factor, + ) + self._coil_clearance = CoilClearanceConstraint( + wall=wall, + coils=coils, + min_clearance_m=float(min_clearance_m), + ) + return self + def _objective(self, x: np.ndarray) -> float: """Objective: minimize total flux on all ports, penalize constraints.""" total = 0.0 @@ -812,6 +851,9 @@ def _objective(self, x: np.ndarray) -> float: # Penalty for port overlapping exclusion zones (check all corners + center) if self._port_overlaps_exclusion(theta_c, zeta_c, theta_w, zeta_w): penalty += 1e6 + if self._coil_clearance is not None: + if self._coil_clearance.port_violates(theta_c, zeta_c, theta_w, zeta_w): + penalty += 1e6 # Penalty for exceeding max flux constraint if self._max_flux_constraint is not None: @@ -1043,6 +1085,69 @@ def plot_heat_flux_2d( return ax +def plot_heat_flux_with_coil_clearance( + heat_map: WallHeatMap, + *, + coil_file: str | Path, + chartmap_file: str | Path, + min_clearance_m: float, + ax=None, + cmap: str = "hot", + vmax: Optional[float] = None, + show_colorbar: bool = True, + ports: Optional[list[Port]] = None, + forbidden_alpha: float = 0.35, +): + """ + Plot heat flux with an overlaid forbidden region based on coil clearance. + + The forbidden region is where the minimum distance from the wall to any + coil segment is less than min_clearance_m. + """ + import matplotlib.pyplot as plt + + if min_clearance_m <= 0.0: + raise ValueError("min_clearance_m must be positive") + + if ax is None: + fig, ax = plt.subplots(figsize=(12, 6)) + + ax = plot_heat_flux_2d( + heat_map, + ax=ax, + cmap=cmap, + vmax=vmax, + show_colorbar=show_colorbar, + ports=ports, + ) + + wall = WallSurface.from_chartmap(chartmap_file) + coils = CoilSegments.from_file(coil_file) + clearance = clearance_map_on_heatmap_grid( + wall, + coils, + heat_map.theta_centers, + heat_map.zeta_centers, + ) + forbidden = clearance < min_clearance_m + + forbidden_grid = forbidden.astype(float) + ax.pcolormesh( + np.degrees(heat_map.zeta_edges), + np.degrees(heat_map.theta_edges), + forbidden_grid, + cmap="Greys", + vmin=0.0, + vmax=1.0, + shading="flat", + alpha=forbidden_alpha, + ) + ax.set_title( + f"Wall Heat Flux + Coil Clearance (d > {min_clearance_m:.3f} m)" + ) + return ax + + def plot_heat_flux_3d( heat_map: WallHeatMap, chartmap_file: Optional[str | Path] = None, diff --git a/test/python/CMakeLists.txt b/test/python/CMakeLists.txt index 2d725780..05c67741 100644 --- a/test/python/CMakeLists.txt +++ b/test/python/CMakeLists.txt @@ -114,7 +114,7 @@ add_test(NAME test_engineering COMMAND ${Python_EXECUTABLE} -m pytest ${CMAKE_CURRENT_SOURCE_DIR}/test_engineering.py -v ) set_tests_properties(test_engineering PROPERTIES - ENVIRONMENT "PYTHONPATH=${CMAKE_SOURCE_DIR}/python" + ENVIRONMENT "PYTHONPATH=${CMAKE_SOURCE_DIR}/python;SIMPLE_TEST_ARTIFACTS_DIR=${CMAKE_CURRENT_BINARY_DIR}/test_artifacts" LABELS "python;engineering" TIMEOUT 300 WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR} diff --git a/test/python/test_engineering.py b/test/python/test_engineering.py index b2467a44..8ef30b91 100644 --- a/test/python/test_engineering.py +++ b/test/python/test_engineering.py @@ -29,6 +29,7 @@ PortOptimizer, OptimizationResult, plot_heat_flux_2d, + plot_heat_flux_with_coil_clearance, compute_wall_area_from_chartmap, ) @@ -657,6 +658,31 @@ def create_mock_chartmap_nc( v_nfp.assignValue(nfp) +def create_mock_coil_file( + path: Path, + *, + R0: float = 10.0, + a: float = 1.0, + clearance: float = 0.05, + n_points: int = 64, +): + """ + Create a simple coil polyline file near the outboard midplane. + + The coil is a vertical line segment at zeta=0 (x axis), positioned just + outside the wall at theta=0, so it constrains ports near zeta=0. + """ + x = np.full((n_points,), R0 + a + clearance, dtype=float) + y = np.zeros((n_points,), dtype=float) + z = np.linspace(-1.5 * a, 1.5 * a, n_points, dtype=float) + I = np.full((n_points,), 1.0, dtype=float) + + data = np.column_stack([x, y, z, I]) + with path.open("w", encoding="utf-8") as f: + f.write(f"{n_points}\n") + np.savetxt(f, data, fmt="%.10e") + + @pytest.mark.skipif(not HAS_NETCDF4, reason="netCDF4 not available") class TestWallAreaCalculation: """Tests for wall area calculation from chartmap.""" @@ -676,6 +702,78 @@ def test_simple_torus_area(self, tmp_path: Path): # Should match within 1.5% (numerical discretization error) assert area == pytest.approx(expected_area, rel=0.015) + +@pytest.mark.skipif(not HAS_NETCDF4, reason="netCDF4 not available") +class TestCoilClearance: + """Tests for coil clearance constraint and plotting.""" + + def test_coil_clearance_penalizes_forbidden_ports( + self, tmp_path: Path, mock_results_file: Path + ): + heat_map = WallHeatMap.from_netcdf( + mock_results_file, total_alpha_power_MW=TEST_ALPHA_POWER_MW + ) + + chartmap_path = tmp_path / "chartmap_nfp1.nc" + create_mock_chartmap_nc(chartmap_path, R0=10.0, a=1.0, nfp=1, nzeta=64) + + coil_path = tmp_path / "coil.dat" + create_mock_coil_file(coil_path, R0=10.0, a=1.0, clearance=0.05) + + opt = PortOptimizer(heat_map) + opt.add_port("P", theta_width=0.2, zeta_width=0.2) + opt.set_coil_clearance_constraint( + coil_path, + chartmap_file=chartmap_path, + min_clearance_m=0.1, + ) + + # zeta=0 should be too close to coil, zeta=pi should be far. + bad = opt._objective(np.array([0.0, 0.0])) + good = opt._objective(np.array([0.0, np.pi])) + assert bad > good + 1e5 + + def test_plot_heat_flux_with_coil_clearance_writes_artifact( + self, tmp_path: Path, mock_results_file: Path + ): + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + heat_map = WallHeatMap.from_netcdf( + mock_results_file, total_alpha_power_MW=TEST_ALPHA_POWER_MW + ) + + chartmap_path = tmp_path / "chartmap_nfp1.nc" + create_mock_chartmap_nc(chartmap_path, R0=10.0, a=1.0, nfp=1, nzeta=64) + + coil_path = tmp_path / "coil.dat" + create_mock_coil_file(coil_path, R0=10.0, a=1.0, clearance=0.05) + + artifacts_root = Path( + os.environ.get("SIMPLE_TEST_ARTIFACTS_DIR", str(tmp_path)) + ) + out_dir = artifacts_root / "engineering" + out_dir.mkdir(parents=True, exist_ok=True) + out_png = out_dir / "port_feasibility_coil_clearance.png" + + fig, ax = plt.subplots(figsize=(12, 6)) + ports = [Port("P", 0.0, np.pi, 0.2, 0.2)] + plot_heat_flux_with_coil_clearance( + heat_map, + coil_file=coil_path, + chartmap_file=chartmap_path, + min_clearance_m=0.1, + ax=ax, + ports=ports, + ) + fig.savefig(out_png, dpi=150) + plt.close(fig) + + assert out_png.exists() + assert out_png.stat().st_size > 10_000 + def test_area_positive(self, tmp_path: Path): """Test that computed area is positive and reasonable.""" chartmap_path = tmp_path / "chartmap.nc" From 9e4ce5820c44ba9d2363a7ad3c98934d931b0c0f Mon Sep 17 00:00:00 2001 From: Christopher Albert Date: Mon, 5 Jan 2026 23:18:20 +0100 Subject: [PATCH 2/2] test: parallelize chartmap wall losses --- .../tests/magfie/test_chartmap_wall_losses.f90 | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/test/tests/magfie/test_chartmap_wall_losses.f90 b/test/tests/magfie/test_chartmap_wall_losses.f90 index bb0f1fc4..8bb82659 100644 --- a/test/tests/magfie/test_chartmap_wall_losses.f90 +++ b/test/tests/magfie/test_chartmap_wall_losses.f90 @@ -19,6 +19,7 @@ program test_chartmap_wall_losses !> Outputs NetCDF file for visualization of loss positions and times. use, intrinsic :: iso_fortran_env, only: dp => real64 + use omp_lib, only: omp_get_max_threads use netcdf use simple, only: init_vmec use util, only: twopi @@ -89,6 +90,7 @@ program test_chartmap_wall_losses print *, '========================================================' print *, 'Comparing loss fractions: VMEC vs Chartmap vs Meiss' print *, '========================================================' + print '(A,I0)', ' OpenMP max threads: ', omp_get_max_threads() print * chartmap_file = 'wout.chartmap.nc' @@ -348,6 +350,7 @@ subroutine trace_particles_vmec(z0, loss_times, loss_pos, lost_step) loss_pos = 0.0_dp lost_step = n_steps +!$omp parallel do default(shared) private(i,j,ierr,z) schedule(static) do i = 1, n_particles z = z0(:, i) do j = 1, n_steps @@ -360,6 +363,7 @@ subroutine trace_particles_vmec(z0, loss_times, loss_pos, lost_step) end if end do end do +!$omp end parallel do end subroutine trace_particles_vmec subroutine trace_particles_chartmap(z0, loss_times, loss_pos, lost_step) @@ -375,6 +379,7 @@ subroutine trace_particles_chartmap(z0, loss_times, loss_pos, lost_step) loss_pos = 0.0_dp lost_step = n_steps +!$omp parallel do default(shared) private(i,j,ierr,z,z_chart) schedule(static) do i = 1, n_particles call vmec_to_chartmap(z0(:, i), z_chart) z = z_chart @@ -388,6 +393,7 @@ subroutine trace_particles_chartmap(z0, loss_times, loss_pos, lost_step) end if end do end do +!$omp end parallel do end subroutine trace_particles_chartmap subroutine trace_particles_boozer(z0, loss_times, loss_pos, lost_step) @@ -406,6 +412,7 @@ subroutine trace_particles_boozer(z0, loss_times, loss_pos, lost_step) loss_pos = 0.0_dp lost_step = n_steps +!$omp parallel do default(shared) private(i,j,ierr,z,theta_b,phi_b,theta_v,phi_v) schedule(static) do i = 1, n_particles z = z0(:, i) call vmec_to_boozer(z(1), z(2), z(3), theta_b, phi_b) @@ -424,6 +431,7 @@ subroutine trace_particles_boozer(z0, loss_times, loss_pos, lost_step) end if end do end do +!$omp end parallel do end subroutine trace_particles_boozer subroutine trace_particles_meiss(z0, loss_times, loss_pos, lost_step) @@ -444,6 +452,7 @@ subroutine trace_particles_meiss(z0, loss_times, loss_pos, lost_step) loss_pos = 0.0_dp lost_step = n_steps +!$omp parallel do default(shared) private(i,j,ierr,z,z_chart,x_ref,x_integ) schedule(static) do i = 1, n_particles call vmec_to_chartmap(z0(:, i), z_chart) x_ref = z_chart(1:3) @@ -463,6 +472,7 @@ subroutine trace_particles_meiss(z0, loss_times, loss_pos, lost_step) end if end do end do +!$omp end parallel do end subroutine trace_particles_meiss subroutine trace_particles_meiss_vmec(z0, loss_times, loss_pos, lost_step) @@ -482,6 +492,7 @@ subroutine trace_particles_meiss_vmec(z0, loss_times, loss_pos, lost_step) loss_pos = 0.0_dp lost_step = n_steps +!$omp parallel do default(shared) private(i,j,ierr,z,x_ref,x_integ) schedule(static) do i = 1, n_particles x_ref = z0(1:3, i) call ref_to_integ(x_ref, x_integ) @@ -500,6 +511,7 @@ subroutine trace_particles_meiss_vmec(z0, loss_times, loss_pos, lost_step) end if end do end do +!$omp end parallel do end subroutine trace_particles_meiss_vmec subroutine trace_particles_boozer_sympl(z0, loss_times, loss_pos, lost_step) @@ -523,6 +535,7 @@ subroutine trace_particles_boozer_sympl(z0, loss_times, loss_pos, lost_step) dtau_sympl = dtau/sqrt(2.0_dp) +!$omp parallel do default(shared) private(i,j,ierr,f,integ,z,z5,theta_b,phi_b,theta_v,phi_v) schedule(static) do i = 1, n_particles z5 = z0(:, i) call vmec_to_boozer(z5(1), z5(2), z5(3), theta_b, phi_b) @@ -552,6 +565,7 @@ subroutine trace_particles_boozer_sympl(z0, loss_times, loss_pos, lost_step) end if end do end do +!$omp end parallel do end subroutine trace_particles_boozer_sympl subroutine trace_particles_meiss_vmec_sympl(z0, loss_times, loss_pos, lost_step) @@ -581,6 +595,7 @@ subroutine trace_particles_meiss_vmec_sympl(z0, loss_times, loss_pos, lost_step) ntau_substeps = max(1, ceiling(dtau/dtaumin)) dtau_sympl = dtaumin/sqrt(2.0_dp) +!$omp parallel do default(shared) private(i,j,ierr,f,integ,z,z5,x_ref,x_integ) schedule(static) do i = 1, n_particles x_ref = z0(1:3, i) call ref_to_integ(x_ref, x_integ) @@ -608,6 +623,7 @@ subroutine trace_particles_meiss_vmec_sympl(z0, loss_times, loss_pos, lost_step) end if end do end do +!$omp end parallel do end subroutine trace_particles_meiss_vmec_sympl subroutine trace_particles_meiss_chart_sympl(z0, loss_times, loss_pos, lost_step) @@ -637,6 +653,7 @@ subroutine trace_particles_meiss_chart_sympl(z0, loss_times, loss_pos, lost_step ntau_substeps = max(1, ceiling(dtau/dtaumin)) dtau_sympl = dtaumin/sqrt(2.0_dp) +!$omp parallel do default(shared) private(i,j,ierr,f,integ,z,z5,z_chart,x_ref,x_integ) schedule(static) do i = 1, n_particles call vmec_to_chartmap(z0(:, i), z_chart) x_ref = z_chart(1:3) @@ -665,6 +682,7 @@ subroutine trace_particles_meiss_chart_sympl(z0, loss_times, loss_pos, lost_step end if end do end do +!$omp end parallel do end subroutine trace_particles_meiss_chart_sympl subroutine vmec_to_chartmap(z_vmec, z_chart)