Skip to content

Commit

Permalink
added new tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@lgiacome
lgiacome committed Nov 18, 2024
1 parent 88973b9 commit 71a5f6c
Showing 1 changed file with 223 additions and 3 deletions.
226 changes: 223 additions & 3 deletions tests/test_utilities.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,19 @@
create_resistive_wall_single_layer_approx_component,
create_resistive_wall_single_layer_approx_element,
create_taper_RW_approx_component,
create_taper_RW_approx_element)
create_element_from_table,
create_interpolated_component)

from xwakes.wit.component import ComponentSingleLayerResistiveWall
from xwakes.wit.interface_dataclasses import _IW2DInputBase
from xwakes.wit.component import KIND_DEFINITIONS

import xobjects as xo

from test_common import relative_error
from pywit.parameters import *
from pywit.interface import (FlatIW2DInput, RoundIW2DInput, Sampling,
component_names)
from xwakes.wit.interface_dataclasses import _IW2DInputBase
from pywit.materials import layer_from_json_material_library, copper_at_temperature

from typing import Dict
Expand All @@ -18,7 +24,7 @@
from pytest import raises, mark, fixture
import numpy as np
from numpy import testing as npt

import pandas as pd

def test_incompatible_dictionaries():
rs = {"z0000": 2}
Expand Down Expand Up @@ -460,3 +466,217 @@ def test_taper_RW_algorithm(freq, component_id, round_single_layer_input):
round_single_layer_input.length/float(npts))

npt.assert_allclose(comp_taper_trapz.impedance(freq), z_steps, rtol=1e-3)

def test_element_from_wake_table():
wake_table = pd.DataFrame.from_dict({
'time': np.linspace(0, 1, 100),
'longitudinal': np.random.rand(100),
'dipolar_x': np.random.rand(100),
'dipolar_y': np.random.rand(100),
'quadrupolar_x': np.random.rand(100),
'quadrupolar_y': np.random.rand(100)
})

test_points = np.linspace(np.min(wake_table['time']),
np.max(wake_table['time']),
100)

element = create_element_from_table(wake_table=wake_table,
use_components=wake_table.columns[1:].to_list(),
length=1,
beta_x=1,
beta_y=1,
name="test element",
tag="Tag",
description='description')

for component in element.components:
for component_kind in KIND_DEFINITIONS:
kind_definition = KIND_DEFINITIONS[component_kind]
component_found = False
if (kind_definition['plane'] == component.plane and
kind_definition['test_exponents'] == component.test_exponents and
kind_definition['source_exponents'] == component.source_exponents):
component_found = True
break

assert component_found

test_wake = np.interp(test_points, wake_table['time'], wake_table[component_kind])

xo.assert_allclose(component.wake(test_points), test_wake, rtol=1e-3)

def test_element_from_impedance_table():
impedance_table = pd.DataFrame.from_dict({
'frequency': np.linspace(1e9, 1e10, 100),
'longitudinal': np.random.rand(100),
'dipolar_x': np.random.rand(100),
'dipolar_y': np.random.rand(100),
'quadrupolar_x': np.random.rand(100),
'quadrupolar_y': np.random.rand(100)
})

test_points = np.linspace(np.min(impedance_table['frequency']),
np.max(impedance_table['frequency']),
100)

element = create_element_from_table(impedance_table=impedance_table,
use_components=impedance_table.columns[1:].to_list(),
length=1,
beta_x=1,
beta_y=1,
name="test element",
tag="Tag",
description='description')

for component in element.components:
for component_kind in KIND_DEFINITIONS:
kind_definition = KIND_DEFINITIONS[component_kind]
component_found = False
if (kind_definition['plane'] == component.plane and
kind_definition['test_exponents'] == component.test_exponents and
kind_definition['source_exponents'] == component.source_exponents):
component_found = True
break

assert component_found

test_impedance = np.interp(test_points, impedance_table['frequency'],
impedance_table[component_kind])

xo.assert_allclose(component.impedance(test_points), test_impedance, rtol=1e-3)

def test_element_from_wake_and_impedance_table():
wake_table = pd.DataFrame.from_dict({
'time': np.linspace(0, 1, 100),
'longitudinal': np.random.rand(100),
'dipolar_x': np.random.rand(100),
'dipolar_y': np.random.rand(100),
'quadrupolar_x': np.random.rand(100),
'quadrupolar_y': np.random.rand(100)
})

impedance_table = pd.DataFrame.from_dict({
'frequency': np.linspace(1e9, 1e10, 100),
'longitudinal': np.random.rand(100),
'dipolar_x': np.random.rand(100),
'dipolar_y': np.random.rand(100),
'quadrupolar_x': np.random.rand(100),
'quadrupolar_y': np.random.rand(100)
})

test_points_wake = np.linspace(np.min(wake_table['time']),
np.max(wake_table['time']),
100)

test_points_impedance = np.linspace(np.min(impedance_table['frequency']),
np.max(impedance_table['frequency']),
100)

element = create_element_from_table(wake_table=wake_table,
impedance_table=impedance_table,
use_components=wake_table.columns[1:].to_list(),
length=1,
beta_x=1,
beta_y=1,
name="test element",
tag="Tag",
description='description')

for component in element.components:
for component_kind in KIND_DEFINITIONS:
kind_definition = KIND_DEFINITIONS[component_kind]
component_found = False
if (kind_definition['plane'] == component.plane and
kind_definition['test_exponents'] == component.test_exponents and
kind_definition['source_exponents'] == component.source_exponents):
component_found = True
break

assert component_found

if hasattr(component, 'wake'):
test_wake = np.interp(test_points_wake, wake_table['time'],
wake_table[component_kind])

xo.assert_allclose(component.wake(test_points_wake), test_wake,
rtol=1e-3)

if hasattr(component, 'impedance'):
test_impedance = np.interp(test_points_impedance,
impedance_table['frequency'],
impedance_table[component_kind])

xo.assert_allclose(component.impedance(test_points_impedance),
test_impedance, rtol=1e-3)


def test_element_from_wake_table_wrong_columns():
wake_table = pd.DataFrame.from_dict({
'time': np.linspace(0, 1, 100),
'longitudinal': np.random.rand(100),
'dipolar_x': np.random.rand(100),
'dipolar_y': np.random.rand(100),
'quadrupolar_x': np.random.rand(100),
'quadrupolar_y': np.random.rand(100)
})

with raises(ValueError):
element = create_element_from_table(wake_table=wake_table,
use_components=['wrong_column'],
length=1,
beta_x=1,
beta_y=1,
name="test element",
tag="Tag",
description='description')

def test_element_from_impedance_table_wrong_columns():
impedance_table = pd.DataFrame.from_dict({
'frequency': np.linspace(1e9, 1e10, 100),
'longitudinal': np.random.rand(100),
'dipolar_x': np.random.rand(100),
'dipolar_y': np.random.rand(100),
'quadrupolar_x': np.random.rand(100),
'quadrupolar_y': np.random.rand(100)
})

with raises(ValueError):
element = create_element_from_table(impedance_table=impedance_table,
use_components=['wrong_column'],
length=1,
beta_x=1,
beta_y=1,
name="test element",
tag="Tag",
description='description')

def test_interpolated_component():
wake_callable = lambda t: 2*t
impedance_callable = lambda f: 3*f

interpolation_frequencies = np.linspace(0, 1, 10)
interpolation_times = np.linspace(10, 100, 10)


component = create_interpolated_component(
plane='z',
source_exponents=(0, 0),
test_exponents=(0, 0),
wake=wake_callable,
impedance=impedance_callable,
interpolation_frequencies=interpolation_frequencies,
interpolation_times=interpolation_times)

test_points_wake = np.linspace(np.min(interpolation_times),
np.max(interpolation_times),
100)

test_points_impedance = np.linspace(np.min(interpolation_frequencies),
np.max(interpolation_frequencies),
100)

xo.assert_allclose(component.wake(test_points_wake),
wake_callable(test_points_wake))
xo.assert_allclose(component.impedance(test_points_impedance),
impedance_callable(test_points_impedance))

0 comments on commit 71a5f6c

Please sign in to comment.