ChristopherMayes / lume-astra

Python wrapper for Astra (A Space Charge Tracking Algorithm, DESY) for eventual use in LUME https://christophermayes.github.io/lume-astra/
https://christophermayes.github.io/lume-astra/
Apache License 2.0
14 stars 9 forks source link

accessing Generator parameters from Astra object or HDF5 archive file #19

Closed kathryn-baker closed 1 year ago

kathryn-baker commented 1 year ago

Hello,

I have run a number of simulations using the run_astra_with_generator function, passing a dictionary of settings that includes settings for both the Astra object and the Generator object.

I saved the output of those simulations as HDF5 files using the to_hdf5 method. I am able to find some of the settings from the dictionary that form part of the Astra object (the cavities and quadrupole strengths for example) in the HDF5 file and resulting Astra object but I can't find the settings for the Generator. Is there any way of accessing these parameters (e.g. sig_x, q_total) from the HDF5 file or the loaded Astra object from that file?

Thanks for your help!

ChristopherMayes commented 1 year ago

Okay, I added archiving in #20. @kathryn-baker could you test that branch?

kathryn-baker commented 1 year ago

Thank you!

I have tried installing that branch using

pip install git+https://github.com/ChristopherMayes/lume-astra.git@generator_archiving.

The installation works so I see the following in my pip list:

lume-astra                    0.5.1+10.g1fbec9a
lume-base                     0.2.3
openpmd-beamphysics           0.7.2

However, when I try to import the Astra object,

from astra import Astra

I get an import error:

ImportError                               Traceback (most recent call last)
Cell In [3], line 1
----> 1 from astra import Astra

File /opt/conda/lib/python3.9/site-packages/astra/__init__.py:4
      2 from .generator import AstraGenerator
      3 from .evaluate import evaluate_astra_with_generator
----> 4 from .astra_distgen import run_astra_with_distgen, evaluate_astra_with_distgen
      6 from . import _version
      7 __version__ = _version.get_versions()['version']

File /opt/conda/lib/python3.9/site-packages/astra/astra_distgen.py:9
      6 from .astra import recommended_spacecharge_mesh
      7 from .evaluate import default_astra_merit
----> 9 from distgen import Generator   
     10 from distgen.writers import write_astra
     11 from distgen.tools import update_nested_dict

ImportError: cannot import name 'Generator' from 'distgen' (/opt/conda/lib/python3.9/site-packages/distgen/__init__.py)
ChristopherMayes commented 1 year ago

Unfortunately distgen on PyPI is a completely different package: https://github.com/devexp-db/distgen. You have to install manually: https://colwyngulliford.github.io/distgen/installation/

kathryn-baker commented 1 year ago

I've managed to resolve the environment and have installed the correct distgen so everything seems to be working as expected in terms of functionality.

However, when I run the following code, the .h5 file still doesn't show the generator input that I can see from the changes you've added.

from astra import run_astra_with_generator
from astra.install import install_astra, install_generator

astra_bin = install_astra(".", verbose=True)
generator_bin = install_generator(".", verbose=True)

settings = {"ipart": 200, "sig_x": 0.1}

test = run_astra_with_generator(
    settings=settings,
    astra_input_file="astra.in",
    generator_input_file="generator.in",
    workdir=".",
    command=astra_bin,
    command_generator=generator_bin,
    timeout=2500,
    verbose=True,
    auto_set_spacecharge_mesh=True,
)

test.archive("test.h5")

Should I be using the evaluate_astra_with_generator function instead of run_astra_with_generator?

Thanks again for your help.

ChristopherMayes commented 1 year ago

evaluate_astra_with_generator does this automatically. You can see what it does: https://github.com/ChristopherMayes/lume-astra/blob/1fbec9a6eea3773c05b5729c4c4e490d14a11e13/astra/evaluate.py#L211

But I think I can also add some of this to run_astra_with_generator.

kathryn-baker commented 1 year ago

Sorry, I seem to be encountering more issues now while trying to test the implementation. I am running the following code with the files I uploaded as part of issue #14.

from astra import run_astra_with_generator, Astra, AstraGenerator
from astra.evaluate import evaluate_astra_with_generator
from astra.install import install_astra, install_generator
from pprint import pprint

astra_bin = install_astra(".", verbose=True)
generator_bin = install_generator(".", verbose=True)

settings = {"ipart": 200, "sig_x": 0.1}

result = evaluate_astra_with_generator(
    settings=settings,
    astra_input_file="ISISMEBT.in",
    generator_input_file="generator_ISISMEBT.in",
    workdir=".",
    astra_bin=astra_bin,
    generator_bin=generator_bin,
    timeout=2500,
    verbose=True,
    # auto_set_spacecharge_mesh=True,
    archive_path="./",
)
pprint(result)

astra_object = Astra.from_archive(result["archive"])
gen_object = AstraGenerator.from_archive(result["archive"])

However, I am getting an index error which I can see is because 0 particle files have been written. I'm not sure if this is related to the archiving solution you've implemented but if there's anything you can suggest to fix this issue I can test the archiving.

Initial particles written to /tmp/tmpu2wfpn7s/astra.particles
loading 0 particle files
[]
{'start_time': 1669114996.1606102, 'run_script': '/src/app/Astra ISISMEBT.in', 'why_error': '', 'run_time': 48.60497307777405, 'run_error': False}
run_astra_with_generator finished
Traceback (most recent call last):
  File "test_astra.py", line 27, in <module>
    result = evaluate_astra_with_generator(
  File "/opt/conda/envs/astra_simulations/lib/python3.8/site-packages/astra/evaluate.py", line 206, in evaluate_astra_with_generator      
    output = default_astra_merit(A)
  File "/opt/conda/envs/astra_simulations/lib/python3.8/site-packages/astra/evaluate.py", line 47, in default_astra_merit
    P = A.particles[-1]
IndexError: list index out of range

I am now running Python 3.8.15 in the continuumio/miniconda3 environment from Docker. I have the following packages and versions installed:

# Name                    Version                   Build  Channel
_libgcc_mutex             0.1                 conda_forge    conda-forge
_openmp_mutex             4.5                       2_gnu    conda-forge
brotli                    1.0.9                h166bdaf_8    conda-forge
brotli-bin                1.0.9                h166bdaf_8    conda-forge
bzip2                     1.0.8                h7f98852_4    conda-forge
c-ares                    1.18.1               h7f98852_0    conda-forge
ca-certificates           2022.9.24            ha878542_0    conda-forge
cached-property           1.5.2                hd8ed1ab_1    conda-forge
cached_property           1.5.2              pyha770c72_1    conda-forge
certifi                   2022.9.24          pyhd8ed1ab_0    conda-forge
charset-normalizer        2.1.1                    pypi_0    pypi
contourpy                 1.0.6            py38h43d8883_0    conda-forge
cycler                    0.11.0             pyhd8ed1ab_0    conda-forge
distgen                   0.6.5              pyhd8ed1ab_0    conda-forge
fonttools                 4.38.0           py38h0a891b7_1    conda-forge
freetype                  2.12.1               hca18f0e_0    conda-forge
h5py                      3.7.0           nompi_py38h7927eab_102    conda-forge
hdf5                      1.12.2          nompi_h4df4325_100    conda-forge
idna                      3.4                      pypi_0    pypi
jpeg                      9e                   h166bdaf_2    conda-forge
keyutils                  1.6.1                h166bdaf_0    conda-forge
kiwisolver                1.4.4            py38h43d8883_1    conda-forge
krb5                      1.19.3               h08a2579_0    conda-forge
lcms2                     2.14                 h6ed2654_0    conda-forge
ld_impl_linux-64          2.39                 hcc3a1bd_1    conda-forge
lerc                      4.0.0                h27087fc_0    conda-forge
libblas                   3.9.0           16_linux64_openblas    conda-forge
libbrotlicommon           1.0.9                h166bdaf_8    conda-forge
libbrotlidec              1.0.9                h166bdaf_8    conda-forge
libbrotlienc              1.0.9                h166bdaf_8    conda-forge
libcblas                  3.9.0           16_linux64_openblas    conda-forge
libcurl                   7.86.0               h2283fc2_1    conda-forge
libdeflate                1.14                 h166bdaf_0    conda-forge
libedit                   3.1.20191231         he28a2e2_2    conda-forge
libev                     4.33                 h516909a_1    conda-forge
libffi                    3.4.2                h7f98852_5    conda-forge
libgcc-ng                 12.2.0              h65d4601_19    conda-forge
libgfortran-ng            12.2.0              h69a702a_19    conda-forge
libgfortran5              12.2.0              h337968e_19    conda-forge
libgomp                   12.2.0              h65d4601_19    conda-forge
liblapack                 3.9.0           16_linux64_openblas    conda-forge
libnghttp2                1.47.0               hff17c54_1    conda-forge
libnsl                    2.0.0                h7f98852_0    conda-forge
libopenblas               0.3.21          pthreads_h78a6416_3    conda-forge
libpng                    1.6.39               h753d276_0    conda-forge
libsqlite                 3.40.0               h753d276_0    conda-forge
libssh2                   1.10.0               hf14f497_3    conda-forge
libstdcxx-ng              12.2.0              h46fd767_19    conda-forge
libtiff                   4.4.0                h55922b4_4    conda-forge
libuuid                   2.32.1            h7f98852_1000    conda-forge
libwebp-base              1.2.4                h166bdaf_0    conda-forge
libxcb                    1.13              h7f98852_1004    conda-forge
libzlib                   1.2.13               h166bdaf_4    conda-forge
lume-astra                0.5.1+10.g1fbec9a          pypi_0    pypi
lume-base                 0.2.3              pyhd8ed1ab_0    conda-forge
matplotlib-base           3.6.2            py38hb021067_0    conda-forge
munkres                   1.1.4              pyh9f0ad1d_0    conda-forge
ncurses                   6.3                  h27087fc_1    conda-forge
numpy                     1.23.5           py38h7042d01_0    conda-forge
openjpeg                  2.5.0                h7d73246_1    conda-forge
openpmd-beamphysics       0.7.2              pyhd8ed1ab_0    conda-forge
openssl                   3.0.7                h166bdaf_0    conda-forge
packaging                 21.3               pyhd8ed1ab_0    conda-forge
pillow                    9.2.0            py38h9eb91d8_3    conda-forge
pint                      0.20.1             pyhd8ed1ab_0    conda-forge
pip                       22.3.1             pyhd8ed1ab_0    conda-forge
pthread-stubs             0.4               h36c2ea0_1001    conda-forge
pyparsing                 3.0.9              pyhd8ed1ab_0    conda-forge
python                    3.8.15          h4a9ceb5_0_cpython    conda-forge
python-dateutil           2.8.2              pyhd8ed1ab_0    conda-forge
python_abi                3.8                      3_cp38    conda-forge
pyyaml                    6.0              py38h0a891b7_5    conda-forge
readline                  8.1.2                h0f457ee_0    conda-forge
requests                  2.28.1                   pypi_0    pypi
scipy                     1.9.3            py38h8ce737c_2    conda-forge
setuptools                65.5.1             pyhd8ed1ab_0    conda-forge
six                       1.16.0             pyh6c4a22f_0    conda-forge
tk                        8.6.12               h27826a3_0    conda-forge
unicodedata2              15.0.0           py38h0a891b7_0    conda-forge
urllib3                   1.26.12                  pypi_0    pypi
wheel                     0.38.4             pyhd8ed1ab_0    conda-forge
xorg-libxau               1.0.9                h7f98852_0    conda-forge
xorg-libxdmcp             1.1.3                h7f98852_0    conda-forge
xz                        5.2.6                h166bdaf_0    conda-forge
yaml                      0.2.5                h7f98852_2    conda-forge
zstd                      1.5.2                h6239696_4    conda-forge

Sorry this hasn't been easy to test...

ChristopherMayes commented 1 year ago

Here you have a problem with your astra run. To debug I added:

A =run_astra_with_generator(
    settings=settings,
    astra_input_file="ISISMEBT.in",
    generator_input_file="generator_ISISMEBT.in",
    workdir=".",
    command=astra_bin,
    command_generator=generator_bin,
    timeout=2500,
    verbose=True,
)

which gives IndexError: too many indices for array: array is 1-dimensional, but 2 were indexed. This means that LUME-Astra couldn't parse the output for some reason.

In the terminal, change the directory to A.path and you see that Astra stops early:

 --------------------------------------------------------------------------

     particle position                 z =   0.8578     m
     particle radius is larger than aperture
     Program stops

     particle reaches position         z =   0.8578     m
     time of flight is                 t =    75.47     ns
     final momentum                    p =    434.2     MeV/c
     final phase (cavity 1)      phi_end =    5348.     deg

     Tracking ended with     1 warning(s)

     **********************************************************************

     Emittance information saved to file    : ISISMEBT.Xemit.001            
     Emittance information saved to file    : ISISMEBT.Yemit.001            
     Emittance information saved to file    : ISISMEBT.Zemit.001            
     Lost & found saved to file             : ISISMEBT.LandF.001            
     Ref. part. information saved to file   : ISISMEBT.ref.001              
     Trajectories saved to file             : ISISMEBT.track.001            
     Phase-space distributions logged in    : ISISMEBT.Log.001              
     Core emittance data saved to file      : ISISMEBT.C99emit.001          
     Beam density saved to file             : ISISMEBT.Density.001          
     Trace space emittance saved to file    : ISISMEBT.TRemit.001           
     Sigma matrix saved to file             : ISISMEBT.Sigma.001            

     **********************************************************************

               ATTENTION: PROGRAM IS QUITTING  EARLY !

     finished simulation  Tue Nov 22 08:49:37 2022
     elapsed time   :                 0.2 seconds 
     execution time :                 0.1 seconds 
     system time    :                 0.1 seconds
     Goodbye.
 --------------------------------------------------------------------------

The error messages are hard to understand, so I now added raise ValueError whenever a problem is encountered parsing an output file. In this case: ValueError: ERROR: Empty output file: /var/folders/2f/l5_mybzs30j4qqvyj98w1_nw0000gn/T/tmpru61cyl6/ISISMEBT.Xemit.001

kathryn-baker commented 1 year ago

Hi, thanks for making the change to the error, that does help to clarify!

It turns out one of the problems I was facing was my own stupidity as well... I was trying to pass a bad value for the emittance in the settings which is why the simulation was failing so sorry about that!

Once I managed to set a good emittance, I was able to get the run_astra_with_generator function working, which archives the Astra object as expected with .archive(). As you'd expect, this currently doesn't save the generator related information.

However, I tried running the evaluate_astra_with_generator function with exactly the same parameters (shown below) and I get the following error:

...
Initial particles written to /tmp/tmpguidwst3/astra.particles
/tmp/tmpguidwst3/ISISMEBT.Xemit.001 mean_z
/tmp/tmpguidwst3/ISISMEBT.Xemit.001 mean_t
/tmp/tmpguidwst3/ISISMEBT.Xemit.001 mean_x
/tmp/tmpguidwst3/ISISMEBT.Xemit.001 sigma_x
/tmp/tmpguidwst3/ISISMEBT.Xemit.001 sigma_xp
/tmp/tmpguidwst3/ISISMEBT.Xemit.001 norm_emit_x
/tmp/tmpguidwst3/ISISMEBT.Xemit.001 cov_x__xp/sigma_x
/tmp/tmpguidwst3/ISISMEBT.Yemit.001 mean_z
/tmp/tmpguidwst3/ISISMEBT.Yemit.001 mean_t
/tmp/tmpguidwst3/ISISMEBT.Yemit.001 mean_y
/tmp/tmpguidwst3/ISISMEBT.Yemit.001 sigma_y
/tmp/tmpguidwst3/ISISMEBT.Yemit.001 sigma_yp
/tmp/tmpguidwst3/ISISMEBT.Yemit.001 norm_emit_y
/tmp/tmpguidwst3/ISISMEBT.Yemit.001 cov_y__yp/sigma_y
/tmp/tmpguidwst3/ISISMEBT.Zemit.001 mean_z
/tmp/tmpguidwst3/ISISMEBT.Zemit.001 mean_t
/tmp/tmpguidwst3/ISISMEBT.Zemit.001 mean_kinetic_energy
/tmp/tmpguidwst3/ISISMEBT.Zemit.001 sigma_z
/tmp/tmpguidwst3/ISISMEBT.Zemit.001 sigma_energy
/tmp/tmpguidwst3/ISISMEBT.Zemit.001 norm_emit_z
/tmp/tmpguidwst3/ISISMEBT.Zemit.001 cov_z__energy/sigma_z
/tmp/tmpguidwst3/ISISMEBT.LandF.001 landf_z
/tmp/tmpguidwst3/ISISMEBT.LandF.001 landf_n_particles
/tmp/tmpguidwst3/ISISMEBT.LandF.001 landf_total_charge
/tmp/tmpguidwst3/ISISMEBT.LandF.001 landf_n_lost
/tmp/tmpguidwst3/ISISMEBT.LandF.001 landf_energy_deposited
/tmp/tmpguidwst3/ISISMEBT.LandF.001 landf_energy_exchange
loading 0 particle files
[]
{'start_time': 1669223419.7893748, 'run_script': '/src/app/data/astra_files/Astra ISISMEBT.in', 'why_error': '', 'run_time': 65.17508554458618, 'run_error': False}
run_astra_with_generator finished
Traceback (most recent call last):
  File "/src/app/data/astra_files/test_astra_not_working.py", line 9, in <module>
    result = evaluate_astra_with_generator(
  File "/opt/conda/envs/astra_simulations/lib/python3.9/site-packages/astra/evaluate.py", line 206, in evaluate_astra_with_generator
    output = default_astra_merit(A)
  File "/opt/conda/envs/astra_simulations/lib/python3.9/site-packages/astra/evaluate.py", line 47, in default_astra_merit
    P = A.particles[-1]
IndexError: list index out of range

Code for run_astra_with_generator:

from astra import run_astra_with_generator, Astra, AstraGenerator
from astra.evaluate import evaluate_astra_with_generator
from astra.install import install_astra, install_generator

astra_bin = install_astra(".", verbose=True)
generator_bin = install_generator(".", verbose=True)

test = run_astra_with_generator(
    settings={},
    astra_input_file="ISISMEBT.in",
    generator_input_file="generator_ISISMEBT.in",
    # workdir=".",
    command=astra_bin,
    command_generator=generator_bin,
    # timeout=2500,
    verbose=True,
    # auto_set_spacecharge_mesh=True,
)

test.archive("test.h5")

Code for evaluate_astra_with_generator:

from astra import run_astra_with_generator, Astra, AstraGenerator
from astra.evaluate import evaluate_astra_with_generator
from astra.install import install_astra, install_generator

astra_bin = install_astra(".", verbose=True)
generator_bin = install_generator(".", verbose=True)

result = evaluate_astra_with_generator(
    settings={},
    astra_input_file="ISISMEBT.in",
    generator_input_file="generator_ISISMEBT.in",
    # workdir=".",
    astra_bin=astra_bin,
    generator_bin=generator_bin,
    # timeout=2500,
    verbose=True,
    # auto_set_spacecharge_mesh=True,
)

astra_object = Astra.from_archive(result["archive"])
gen_object = AstraGenerator.from_archive(result["archive"])

I'll be honest I'm not sure what's going on here to cause the difference so this is pure speculation but is it possible that the temp files are causing an issue between the processes? From the error message it seems like the simulation completes because the astra finished message comes through but it subsequently can't find the particles.

ChristopherMayes commented 1 year ago

Could you attach this example?

kathryn-baker commented 1 year ago

This was using the same files as before - I've uploaded them with the python files here astra_files.zip

ChristopherMayes commented 1 year ago

It turns out that evaluate_astra_with_generator was looking for particle files to process, but your example had no output particles written. I fixed this -- could you try now?

kathryn-baker commented 1 year ago

Yes, that's working now! Thank you so much :)