|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | + |
| 3 | +import random |
| 4 | +import shutil |
| 5 | +from dataclasses import dataclass |
| 6 | + |
| 7 | +import numpy as np |
| 8 | +import pytest |
| 9 | + |
| 10 | +import amrex.space3d as amr |
| 11 | + |
| 12 | +# Particle container constructor -- depending on if gpus are available use the |
| 13 | +# GPU-enabled version. This uses the FHDeX Particle Container for the test |
| 14 | +if amr.Config.have_gpu: |
| 15 | + PCC = amr.ParticleContainer_16_4_0_0_managed |
| 16 | +else: |
| 17 | + PCC = amr.ParticleContainer_16_4_0_0_default |
| 18 | + |
| 19 | + |
| 20 | +@dataclass |
| 21 | +class Particle: |
| 22 | + """ |
| 23 | + Helper class to work with particle data |
| 24 | + """ |
| 25 | + |
| 26 | + x: float |
| 27 | + y: float |
| 28 | + z: float |
| 29 | + |
| 30 | + idx: int |
| 31 | + |
| 32 | + def to_soa(self, aos_numpy): |
| 33 | + """Fill amr particle SoA with x, y, z, idx data""" |
| 34 | + aos_numpy["x"] = self.x |
| 35 | + aos_numpy["y"] = self.y |
| 36 | + aos_numpy["z"] = self.z |
| 37 | + aos_numpy["idata_0"] = self.idx |
| 38 | + |
| 39 | + |
| 40 | +def generate_test_particles(n_part): |
| 41 | + """ |
| 42 | + Returns a list of test particles scattered throught the domain |
| 43 | + """ |
| 44 | + particles = list() |
| 45 | + |
| 46 | + def generator(): |
| 47 | + return 1 - 2 * random.random() |
| 48 | + |
| 49 | + for i in range(n_part): |
| 50 | + particles.append(Particle(x=generator(), y=generator(), z=generator(), idx=i)) |
| 51 | + |
| 52 | + return particles |
| 53 | + |
| 54 | + |
| 55 | +def particle_container(Rpart, std_geometry, distmap, boxarr, std_real_box): |
| 56 | + """ |
| 57 | + Generate a fresh particle container, containing copies from Rpart |
| 58 | + """ |
| 59 | + pc = PCC(std_geometry, distmap, boxarr) |
| 60 | + |
| 61 | + iseed = 1 |
| 62 | + myt = amr.ParticleInitType_16_4_0_0() |
| 63 | + pc.init_random(len(Rpart), iseed, myt, False, std_real_box) |
| 64 | + |
| 65 | + particles_tile_ct = 0 |
| 66 | + # assign some values to runtime components |
| 67 | + for lvl in range(pc.finest_level + 1): |
| 68 | + for pti in pc.iterator(pc, level=lvl): |
| 69 | + aos = pti.aos() |
| 70 | + aos_numpy = aos.to_numpy(copy=False) |
| 71 | + for i, p in enumerate(aos_numpy): |
| 72 | + Rpart[i + particles_tile_ct].to_soa(p) |
| 73 | + particles_tile_ct += len(aos_numpy) |
| 74 | + |
| 75 | + pc.redistribute() |
| 76 | + return pc |
| 77 | + |
| 78 | + |
| 79 | +def check_particles_container(pc, reference_particles): |
| 80 | + """ |
| 81 | + Checks the contents of `pc` against `reference_particles` |
| 82 | + """ |
| 83 | + for lvl in range(pc.finest_level + 1): |
| 84 | + for i, pti in enumerate(pc.iterator(pc, level=lvl)): |
| 85 | + aos = pti.aos() |
| 86 | + for p in aos.to_numpy(copy=True): |
| 87 | + ref = reference_particles[p["idata_0"]] |
| 88 | + assert Particle(x=p["x"], y=p["y"], z=p["z"], idx=p["idata_0"]) == ref |
| 89 | + |
| 90 | + |
| 91 | +def write_test_plotfile(filename, reference_part): |
| 92 | + """ |
| 93 | + Write single-level plotfile (in order to read it back in). |
| 94 | + """ |
| 95 | + domain_box = amr.Box([0, 0, 0], [31, 31, 31]) |
| 96 | + real_box = amr.RealBox([-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]) |
| 97 | + geom = amr.Geometry(domain_box, real_box, amr.CoordSys.cartesian, [0, 0, 0]) |
| 98 | + |
| 99 | + ba = amr.BoxArray(domain_box) |
| 100 | + dm = amr.DistributionMapping(ba, 1) |
| 101 | + mf = amr.MultiFab(ba, dm, 1, 0) |
| 102 | + mf.set_val(np.pi) |
| 103 | + |
| 104 | + time = 1.0 |
| 105 | + level_step = 200 |
| 106 | + var_names = amr.Vector_string(["density"]) |
| 107 | + amr.write_single_level_plotfile(filename, mf, var_names, geom, time, level_step) |
| 108 | + |
| 109 | + pc = particle_container(reference_part, geom, dm, ba, real_box) |
| 110 | + pc.write_plotfile(filename, "particles") |
| 111 | + |
| 112 | + |
| 113 | +def load_test_plotfile_particle_container(plot_file_name): |
| 114 | + """ |
| 115 | + Load signle-level plotfile and return particle container |
| 116 | + """ |
| 117 | + plt = amr.PlotFileData(plot_file_name) |
| 118 | + |
| 119 | + probDomain = plt.probDomain(0) |
| 120 | + probLo = plt.probLo() |
| 121 | + probHi = plt.probHi() |
| 122 | + domain_box = amr.Box(probDomain.small_end, probDomain.big_end) |
| 123 | + real_box = amr.RealBox(probLo, probHi) |
| 124 | + std_geometry = amr.Geometry(domain_box, real_box, plt.coordSys(), [0, 0, 0]) |
| 125 | + |
| 126 | + pc = amr.ParticleContainer_16_4_0_0_default( |
| 127 | + std_geometry, |
| 128 | + plt.DistributionMap(plt.finestLevel()), |
| 129 | + plt.boxArray(plt.finestLevel()), |
| 130 | + ) |
| 131 | + pc.restart(plot_file_name, "particles") |
| 132 | + |
| 133 | + return pc |
| 134 | + |
| 135 | + |
| 136 | +@pytest.mark.skipif(amr.Config.spacedim != 3, reason="Requires AMREX_SPACEDIM = 3") |
| 137 | +def test_plotfile_particle_data_read(): |
| 138 | + """ |
| 139 | + Generate and then read a plot file containing particle data. Checks that |
| 140 | + the particle data matches the original particle list. |
| 141 | + """ |
| 142 | + # seed RNG to make test reproducible -- comment out this line to generate new |
| 143 | + # random particle positions every time. |
| 144 | + random.seed(1) |
| 145 | + |
| 146 | + plt_file_name = "plt_test" |
| 147 | + # Reference particle lists |
| 148 | + n_part = 15 |
| 149 | + reference_part = generate_test_particles(n_part) |
| 150 | + # Write a test plotfile containing the reference particles in a paritcle |
| 151 | + # container |
| 152 | + write_test_plotfile(plt_file_name, reference_part) |
| 153 | + # Load the particle container from the test plot file |
| 154 | + pc = load_test_plotfile_particle_container(plt_file_name) |
| 155 | + # Check that the particles in the loaded particle container match the |
| 156 | + # original particle list |
| 157 | + check_particles_container(pc, reference_part) |
| 158 | + |
| 159 | + # clean up after yourself |
| 160 | + shutil.rmtree(plt_file_name) |
0 commit comments