-
Notifications
You must be signed in to change notification settings - Fork 19
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add Particle I/O: restart & restart_checkpoint #342
Merged
Merged
Changes from 22 commits
Commits
Show all changes
23 commits
Select commit
Hold shift + click to select a range
f971b93
add particle io
JBlaschke 5524a84
demo list
JBlaschke 7815c68
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] 37aaad1
fix copy vs alias
JBlaschke 26ee054
Merge branch 'jpb/particle-io' of github.com:JBlaschke/pyamrex into j…
JBlaschke 12c3876
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] aa7df81
Merge branch 'development' into jpb/particle-io
JBlaschke cb17ee3
Merge branch 'jpb/particle-io' of github.com:JBlaschke/pyamrex into j…
JBlaschke 5eb6f30
nearly finished test
JBlaschke c556d24
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] fa20a6f
guard against double intialization
JBlaschke 04c0142
guard against initializing pymarex during pytest
JBlaschke da55bec
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] 6daa936
Add redistribute
ax3l 1978073
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] 4389ba7
Cleanup: double return
ax3l 422e041
Ruff: Lambda -> Def
ax3l ba8818f
Simplify
ax3l 5e00b2c
Alphabetic Order
ax3l b4ca3c2
Alphabetic Order
ax3l 219dfe0
Copyright Header
ax3l e80f90d
Cleanup Test
ax3l 69b36cc
ICC: -j2
ax3l File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,16 @@ | ||
/* Copyright 2024 The AMReX Community | ||
* | ||
* Authors: Johannes Blaschke | ||
* License: BSD-3-Clause-LBNL | ||
*/ | ||
#include "ParticleContainer.H" | ||
|
||
#include <AMReX_Particle.H> | ||
#include <AMReX_ParticleTile.H> | ||
|
||
|
||
void init_ParticleContainer_FHDeX(py::module& m) { | ||
using namespace amrex; | ||
|
||
make_ParticleContainer_and_Iterators<Particle<16, 4>, 0, 0>(m); | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,160 @@ | ||
# -*- coding: utf-8 -*- | ||
|
||
import random | ||
import shutil | ||
from dataclasses import dataclass | ||
|
||
import numpy as np | ||
import pytest | ||
|
||
import amrex.space3d as amr | ||
|
||
# Particle container constructor -- depending on if gpus are available use the | ||
# GPU-enabled version. This uses the FHDeX Particle Container for the test | ||
if amr.Config.have_gpu: | ||
PCC = amr.ParticleContainer_16_4_0_0_managed | ||
else: | ||
PCC = amr.ParticleContainer_16_4_0_0_default | ||
|
||
|
||
@dataclass | ||
class Particle: | ||
""" | ||
Helper class to work with particle data | ||
""" | ||
|
||
x: float | ||
y: float | ||
z: float | ||
|
||
idx: int | ||
|
||
def to_soa(self, aos_numpy): | ||
"""Fill amr particle SoA with x, y, z, idx data""" | ||
aos_numpy["x"] = self.x | ||
aos_numpy["y"] = self.y | ||
aos_numpy["z"] = self.z | ||
aos_numpy["idata_0"] = self.idx | ||
|
||
|
||
def generate_test_particles(n_part): | ||
""" | ||
Returns a list of test particles scattered throught the domain | ||
""" | ||
particles = list() | ||
|
||
def generator(): | ||
return 1 - 2 * random.random() | ||
|
||
for i in range(n_part): | ||
particles.append(Particle(x=generator(), y=generator(), z=generator(), idx=i)) | ||
|
||
return particles | ||
|
||
|
||
def particle_container(Rpart, std_geometry, distmap, boxarr, std_real_box): | ||
""" | ||
Generate a fresh particle container, containing copies from Rpart | ||
""" | ||
pc = PCC(std_geometry, distmap, boxarr) | ||
|
||
iseed = 1 | ||
myt = amr.ParticleInitType_16_4_0_0() | ||
pc.init_random(len(Rpart), iseed, myt, False, std_real_box) | ||
|
||
particles_tile_ct = 0 | ||
# assign some values to runtime components | ||
for lvl in range(pc.finest_level + 1): | ||
for pti in pc.iterator(pc, level=lvl): | ||
aos = pti.aos() | ||
aos_numpy = aos.to_numpy(copy=False) | ||
for i, p in enumerate(aos_numpy): | ||
Rpart[i + particles_tile_ct].to_soa(p) | ||
particles_tile_ct += len(aos_numpy) | ||
ax3l marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
||
pc.redistribute() | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. @JBlaschke This fixed the |
||
return pc | ||
|
||
|
||
def check_particles_container(pc, reference_particles): | ||
""" | ||
Checks the contents of `pc` against `reference_particles` | ||
""" | ||
for lvl in range(pc.finest_level + 1): | ||
for i, pti in enumerate(pc.iterator(pc, level=lvl)): | ||
aos = pti.aos() | ||
for p in aos.to_numpy(copy=True): | ||
ref = reference_particles[p["idata_0"]] | ||
assert Particle(x=p["x"], y=p["y"], z=p["z"], idx=p["idata_0"]) == ref | ||
|
||
|
||
def write_test_plotfile(filename, reference_part): | ||
""" | ||
Write single-level plotfile (in order to read it back in). | ||
""" | ||
domain_box = amr.Box([0, 0, 0], [31, 31, 31]) | ||
real_box = amr.RealBox([-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]) | ||
geom = amr.Geometry(domain_box, real_box, amr.CoordSys.cartesian, [0, 0, 0]) | ||
|
||
ba = amr.BoxArray(domain_box) | ||
dm = amr.DistributionMapping(ba, 1) | ||
mf = amr.MultiFab(ba, dm, 1, 0) | ||
mf.set_val(np.pi) | ||
|
||
time = 1.0 | ||
level_step = 200 | ||
var_names = amr.Vector_string(["density"]) | ||
amr.write_single_level_plotfile(filename, mf, var_names, geom, time, level_step) | ||
|
||
pc = particle_container(reference_part, geom, dm, ba, real_box) | ||
pc.write_plotfile(filename, "particles") | ||
|
||
|
||
def load_test_plotfile_particle_container(plot_file_name): | ||
""" | ||
Load signle-level plotfile and return particle container | ||
""" | ||
plt = amr.PlotFileData(plot_file_name) | ||
|
||
probDomain = plt.probDomain(0) | ||
probLo = plt.probLo() | ||
probHi = plt.probHi() | ||
domain_box = amr.Box(probDomain.small_end, probDomain.big_end) | ||
real_box = amr.RealBox(probLo, probHi) | ||
std_geometry = amr.Geometry(domain_box, real_box, plt.coordSys(), [0, 0, 0]) | ||
|
||
pc = amr.ParticleContainer_16_4_0_0_default( | ||
std_geometry, | ||
plt.DistributionMap(plt.finestLevel()), | ||
plt.boxArray(plt.finestLevel()), | ||
) | ||
pc.restart(plot_file_name, "particles") | ||
|
||
return pc | ||
|
||
|
||
@pytest.mark.skipif(amr.Config.spacedim != 3, reason="Requires AMREX_SPACEDIM = 3") | ||
def test_plotfile_particle_data_read(): | ||
""" | ||
Generate and then read a plot file containing particle data. Checks that | ||
the particle data matches the original particle list. | ||
""" | ||
# seed RNG to make test reproducible -- comment out this line to generate new | ||
# random particle positions every time. | ||
random.seed(1) | ||
|
||
plt_file_name = "plt_test" | ||
# Reference particle lists | ||
n_part = 15 | ||
reference_part = generate_test_particles(n_part) | ||
# Write a test plotfile containing the reference particles in a paritcle | ||
# container | ||
write_test_plotfile(plt_file_name, reference_part) | ||
# Load the particle container from the test plot file | ||
pc = load_test_plotfile_particle_container(plt_file_name) | ||
# Check that the particles in the loaded particle container match the | ||
# original particle list | ||
check_particles_container(pc, reference_part) | ||
|
||
# clean up after yourself | ||
shutil.rmtree(plt_file_name) |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@JBlaschke do you need this for 1D, 2D and 3D AMReX builds or only for 3D?