<jupyter_start><jupyter_text>Scanning Astra evaluations<jupyter_code>from astra import evaluate_astra_with_distgen

import matplotlib.pyplot as plt

import numpy as np
import os
# Input files
ASTRA_IN = 'templates/dcgun/astra.in'
DISTGEN_IN = 'templates/dcgun/distgen.yaml'
# Run Astra with Distgen

# Basic settings
settings0 = {'zstop':1, 'zemit':10, 'zphase':10, 'phases':True}
settings0['distgen:n_particle'] = 2000

# Solenoid field
settings0['maxb(1)'] = 0.04
# This is the basic function we will use
outputs = evaluate_astra_with_distgen(settings0, astra_input_file=ASTRA_IN,
                              distgen_input_file=DISTGEN_IN)
outputs
# Make a bare minimum function to evaluate
def f(b):
    s = settings0.copy()
    s['maxb(1)'] = b
    return evaluate_astra_with_distgen(s, astra_input_file=ASTRA_IN,
                          distgen_input_file=DISTGEN_IN)
# Make a list to scan
blist = np.linspace(0, 0.1, 20)
f(0)
# Get a parallel executor

# Processes. These lines are needed to get this to work on macOS
import multiprocessing as mp                                                                                     
mp.set_start_method("fork")    
from concurrent.futures import ProcessPoolExecutor as Executor

#from concurrent.futures import ThreadPoolExecutor as Executor # DO NOT USE
#from dask.distributed import Client as Executor

executor=Executor()
# Make a map object. Note that this does not execute immediately.
results = executor.map(f, blist)
# Actually evalute the funciton and gather the results
outputs = []
for r in results:
    outputs.append(r)
#    outputs.append(r.result())
# Simple plot
plt.plot(blist, [o['end_sigma_x'] for o in outputs])<jupyter_output><empty_output><jupyter_text>Better: robust evaluation, form dataset<jupyter_code>from tempfile import TemporaryDirectory
SCRATCH=TemporaryDirectory()
# Make a robust evaluate
def evaluate(b):
    try:
        s = settings0.copy()
        s['maxb(1)'] = b
        output = evaluate_astra_with_distgen(s, astra_input_file=ASTRA_IN,
                              distgen_input_file=DISTGEN_IN, archive_path=SCRATCH.name)
        # Add the input
        output['maxb(1)'] = b
        output['Exception'] = False
    except:
        output = {}
        output['maxb(1)'] = b
        output['Exception'] = True
        
    return output
SCRATCH.name
# This will cause a crash in Astra
evaluate(1e9)
# Get a new executor
executor =  Executor()

# Make a map object. Note that this does not execute immediately.

blist[-1] = 1e9 # Add something that will crash
results = executor.map(evaluate, blist)
# Actually evalute the funciton and gather the results
outputs = []
for r in results:
    outputs.append(r)
# Archive files are here
os.listdir(SCRATCH.name)<jupyter_output><empty_output><jupyter_text>Pandas<jupyter_code>import pandas as pd
# Make a DataFrame
df = pd.DataFrame(outputs)
df
# Get all runs without an exeption
filtered_df=df[~df['Exception']]<jupyter_output><empty_output><jupyter_text>Visualize with Bokeh<jupyter_code>from bokeh.plotting import figure, output_notebook, show

from bokeh.models import ColumnDataSource

# Allows plots to show up in the notebook
output_notebook()
# Form a Bokeh style data source
source = ColumnDataSource(filtered_df)
source.column_names
# A simple plot
p = figure()
p.circle(x='maxb(1)', y='end_sigma_x', source=source)
show(p)
# Fancier plot

from bokeh.models.tools import HoverTool

hover = HoverTool()


hover.tooltips=[
    ('sigma_x', '@end_sigma_x'),
    ('norm_emit_x', '@end_norm_emit_x'),
    ('archive', '@archive')
]
p = figure(tools='tap')
p.circle(x='maxb(1)', y='end_sigma_x', source=source, size=10, color='green')

p.add_tools(hover)

show(p)
# Load all archives
from astra import Astra
Alist = []
for a in filtered_df['archive']:
    A = Astra()
    A.load_archive(a)
    Alist.append(A)
plt.plot(
    Alist[0].output['stats']['mean_z'],
    Alist[0].output['stats']['sigma_x']
)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10661302489313552, 'fraction_numerical': 0.010057832537088257, 'mean_word_length': 3.9477611940298507, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '533181', 'license': 'Apache-2.0', 'n_tokens_mistral': 1361, 'n_tokens_neox': 1317, 'n_words': 351, 'path': 'examples/scan_example.ipynb', 'repo_name': 'hhslepicka/lume-astra', 'size': '3977'}