Skip to content

Commit

Permalink
Added a Geometry Observable
Browse files Browse the repository at this point in the history
  • Loading branch information
@lfarv
lfarv committed Jun 3, 2023
1 parent 9ed9e37 commit a4053cb
Show file tree
Hide file tree
Showing 3 changed files with 82 additions and 119 deletions.
131 changes: 21 additions & 110 deletions docs/p/notebooks/test_observables.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@
"outputs": [],
"source": [
"from at import Observable, ObservableList, OrbitObservable, GlobalOpticsObservable, LocalOpticsObservable\n",
"from at import MatrixObservable, TrajectoryObservable, EmittanceObservable, LatticeObservable"
"from at import MatrixObservable, TrajectoryObservable, EmittanceObservable, LatticeObservable, GeometryObservable"
]
},
{
Expand Down Expand Up @@ -457,6 +457,16 @@
"allobs.append(TrajectoryObservable(at.Monitor,axis='px'))"
]
},
{
"cell_type": "code",
"execution_count": 24,
"id": "fef988a4-1413-4478-a8da-8c865b27d74b",
"metadata": {},
"outputs": [],
"source": [
"allobs.append(GeometryObservable(at.Monitor, 'x'))"
]
},
{
"cell_type": "markdown",
"id": "48167f3a-24fa-4963-97ec-1e78033f8895",
Expand All @@ -469,7 +479,7 @@
},
{
"cell_type": "code",
"execution_count": 23,
"execution_count": 25,
"id": "f78c247c-0228-499f-9455-0f39c57ff6fd",
"metadata": {},
"outputs": [],
Expand All @@ -490,7 +500,7 @@
},
{
"cell_type": "code",
"execution_count": 24,
"execution_count": 26,
"id": "582b7317-5527-4afd-ad8c-278f0408dd11",
"metadata": {
"tags": []
Expand All @@ -502,7 +512,7 @@
"array([9.38969042e+00, 2.99742405e+00, 1.40412855e-14])"
]
},
"execution_count": 24,
"execution_count": 26,
"metadata": {},
"output_type": "execute_result"
}
Expand All @@ -521,7 +531,7 @@
},
{
"cell_type": "code",
"execution_count": 25,
"execution_count": 27,
"id": "edfba8ca-ee8f-4064-9076-d2bc469739a4",
"metadata": {},
"outputs": [
Expand Down Expand Up @@ -593,10 +603,12 @@
" 26.374287952316944,\n",
" array([ 0.00000000e+00, -6.94370474e-04, 6.07151649e-04, 2.38468291e-04,\n",
" -6.81824078e-04, -4.78921797e-04, 4.41491589e-04, 7.01582199e-04,\n",
" -6.05543962e-04, -9.78684823e-05])]"
" -6.05543962e-04, -9.78684823e-05]),\n",
" array([ 2.6514 , 6.4783308 , 7.51380991, 10.28830988, 12.71979153,\n",
" 13.62739043, 16.04912109, 18.79203055, 19.81402164, 23.58054559])]"
]
},
"execution_count": 25,
"execution_count": 27,
"metadata": {},
"output_type": "execute_result"
}
Expand All @@ -619,113 +631,12 @@
},
{
"cell_type": "code",
"execution_count": 26,
"execution_count": null,
"id": "6126bf27-37e9-4c20-8997-c82ba9fad5b1",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"location Initial Actual Low bound High bound residual \n",
"orbit[x]\n",
" BPM_01 -3.0219e-09 -3.0219e-09 - - 0.0 \n",
" BPM_02 4.50695e-07 4.50695e-07 - - 0.0 \n",
" BPM_03 4.08206e-07 4.08206e-07 - - 0.0 \n",
" BPM_04 2.379e-08 2.379e-08 - - 0.0 \n",
" BPM_05 -1.31784e-08 -1.31784e-08 - - 0.0 \n",
" BPM_06 2.47231e-08 2.47231e-08 - - 0.0 \n",
" BPM_07 -2.95311e-08 -2.95311e-08 - - 0.0 \n",
" BPM_08 -4.05598e-07 -4.05598e-07 - - 0.0 \n",
" BPM_09 -4.47398e-07 -4.47398e-07 - - 0.0 \n",
" BPM_10 -2.24851e-09 -2.24851e-09 - - 0.0 \n",
"beta[y]\n",
" BPM_01 5.3028 5.3028 -inf 7.0 0.0 \n",
" BPM_02 7.17604 7.17604 -inf 7.0 0.0309906 \n",
" BPM_03 6.55088 6.55088 -inf 7.0 0.0 \n",
" BPM_04 2.31449 2.31449 -inf 7.0 0.0 \n",
" BPM_05 3.40498 3.40498 -inf 7.0 0.0 \n",
" BPM_06 3.40504 3.40504 -inf 7.0 0.0 \n",
" BPM_07 2.31465 2.31465 -inf 7.0 0.0 \n",
" BPM_08 6.55106 6.55106 -inf 7.0 0.0 \n",
" BPM_09 7.17614 7.17614 -inf 7.0 0.0310259 \n",
" BPM_10 5.30284 5.30284 -inf 7.0 0.0 \n",
"matrix\n",
" BPM_02 [-1.082 ...] [-1.082 ...] - - [ 0.0 ...] \n",
"amax(beta[y])\n",
" 7.17614 7.17614 - - 0.0 \n",
"closed_orbit[slice(None, 4, None)]\n",
" QF1A [-3.028e-09 ...] [-3.028e-09 ...] [ 0.0 ...] [ 0.0 ...] [ 9.169e-06 ...] \n",
" QD2A [-1.785e-09 ...] [-1.785e-09 ...] [ 0.0 ...] [ 0.0 ...] [ 3.185e-06 ...] \n",
" QD3A [ 2.06e-07 ...] [ 2.06e-07 ...] [ 0.0 ...] [ 0.0 ...] [ 0.04245 ...] \n",
" QF4A [ 4.635e-07 ...] [ 4.635e-07 ...] [ 0.0 ...] [ 0.0 ...] [ 0.2148 ...] \n",
" QF4B [ 4.928e-07 ...] [ 4.928e-07 ...] [ 0.0 ...] [ 0.0 ...] [ 0.2429 ...] \n",
" QD5B [ 2.391e-07 ...] [ 2.391e-07 ...] [ 0.0 ...] [ 0.0 ...] [ 0.05715 ...] \n",
" QF6B [ 2.249e-08 ...] [ 2.249e-08 ...] [ 0.0 ...] [ 0.0 ...] [ 0.000506 ...] \n",
" QF8B [-2.958e-08 ...] [-2.958e-08 ...] [ 0.0 ...] [ 0.0 ...] [ 0.000875 ...] \n",
" QF8D [ 3.864e-08 ...] [ 3.864e-08 ...] [ 0.0 ...] [ 0.0 ...] [ 0.001493 ...] \n",
" QF6D [-1.147e-08 ...] [-1.147e-08 ...] [ 0.0 ...] [ 0.0 ...] [ 0.0001315 ...] \n",
" QD5D [-1.925e-07 ...] [-1.925e-07 ...] [ 0.0 ...] [ 0.0 ...] [ 0.03705 ...] \n",
" QF4D [-4.585e-07 ...] [-4.585e-07 ...] [ 0.0 ...] [ 0.0 ...] [ 0.2102 ...] \n",
" QF4E [-4.902e-07 ...] [-4.902e-07 ...] [ 0.0 ...] [ 0.0 ...] [ 0.2403 ...] \n",
" QD3E [-2.424e-07 ...] [-2.424e-07 ...] [ 0.0 ...] [ 0.0 ...] [ 0.05877 ...] \n",
" QD2E [-8.05e-10 ...] [-8.05e-10 ...] [ 0.0 ...] [ 0.0 ...] [ 6.48e-07 ...] \n",
" QF1E [-1.927e-09 ...] [-1.927e-09 ...] [ 0.0 ...] [ 0.0 ...] [ 3.715e-06 ...] \n",
"s_pos\n",
" QF1A 2.69395 2.69395 - - 0.0 \n",
" QD2A 3.42956 3.42956 - - 0.0 \n",
" QD3A 5.52309 5.52309 - - 0.0 \n",
" QF4A 6.52741 6.52741 - - 0.0 \n",
" QF4B 7.08941 7.08941 - - 0.0 \n",
" QD5B 8.14327 8.14327 - - 0.0 \n",
" QF6B 10.3428 10.3428 - - 0.0 \n",
" QF8B 11.9398 11.9398 - - 0.0 \n",
" QF8D 13.9418 13.9418 - - 0.0 \n",
" QF6D 15.6329 15.6329 - - 0.0 \n",
" QD5D 18.0033 18.0033 - - 0.0 \n",
" QF4D 19.0572 19.0572 - - 0.0 \n",
" QF4E 19.6192 19.6192 - - 0.0 \n",
" QD3E 20.6726 20.6726 - - 0.0 \n",
" QD2E 22.717 22.717 - - 0.0 \n",
" QF1E 23.3684 23.3684 - - 0.0 \n",
"phase_advance\n",
" [ 9.39 ...] [ 9.39 ...] - - [ 0.0 ...] \n",
"tune[x]\n",
" 2.38156 2.38156 - - 0.0 \n",
"mu\n",
" End [ 14.96 ...] [ 14.96 ...] - - [ 0.0 ...] \n",
"chromaticity\n",
" [ 0.1792 ...] [ 0.1792 ...] - - [ 0.0 ...] \n",
"mean(H)\n",
" -25.3692 -25.3692 - - 0.0 \n",
"PolynomB[2]\n",
" SD1A -78.9554 -78.9554 - - 0.0 \n",
" SF2A 77.0372 77.0372 - - 0.0 \n",
" SD1B -74.1895 -74.1895 - - 0.0 \n",
" SD1D -74.1895 -74.1895 - - 0.0 \n",
" SF2E 77.0372 77.0372 - - 0.0 \n",
" SD1E -78.9554 -78.9554 - - 0.0 \n",
"emittances[x]\n",
" 1.32039e-10 1.32039e-10 - - 0.0 \n",
"circumference\n",
" 26.3743 26.3743 - - 0.0 \n",
"trajectory[px]\n",
" BPM_01 0.0 0.0 - - 0.0 \n",
" BPM_02 -0.00069437 -0.00069437 - - 0.0 \n",
" BPM_03 0.000607152 0.000607152 - - 0.0 \n",
" BPM_04 0.000238468 0.000238468 - - 0.0 \n",
" BPM_05 -0.000681824 -0.000681824 - - 0.0 \n",
" BPM_06 -0.000478922 -0.000478922 - - 0.0 \n",
" BPM_07 0.000441492 0.000441492 - - 0.0 \n",
" BPM_08 0.000701582 0.000701582 - - 0.0 \n",
" BPM_09 -0.000605544 -0.000605544 - - 0.0 \n",
" BPM_10 -9.78685e-05 -9.78685e-05 - - 0.0 \n"
]
}
],
"outputs": [],
"source": [
"print(allobs)"
]
Expand Down
12 changes: 6 additions & 6 deletions pyat/at/lattice/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1065,13 +1065,13 @@ def get_geometry(ring: List[Element],
>>> geomdata, radius = get_geometry(ring)
"""

geom_dtype = [('x', numpy.float64, (1, )),
('y', numpy.float64, (1, )),
('angle', numpy.float64, (1, ))]
geom_dtype = [('x', numpy.float64),
('y', numpy.float64),
('angle', numpy.float64)]
geomdata = numpy.recarray((len(ring)+1, ), dtype=geom_dtype)
xx = numpy.zeros((len(ring)+1, 1))
yy = numpy.zeros((len(ring)+1, 1))
angle = numpy.zeros((len(ring)+1, 1))
xx = numpy.zeros(len(ring)+1)
yy = numpy.zeros(len(ring)+1)
angle = numpy.zeros(len(ring)+1)
x, y, t = start_coordinates
x0, y0, t0 = start_coordinates

Expand Down
58 changes: 55 additions & 3 deletions pyat/at/latticetools/observables.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,9 @@ class Need(Enum):
ALL_POINTS = 7
#: Associated with LOCALOPTICS, require the *get_chrom* keyword
CHROMATICITY = 8
#: Specify geometry computation and provide the full data at evaluation
#: points
GEOMETRY = 9


class Observable(object):
Expand Down Expand Up @@ -447,6 +450,50 @@ def _setup(self, ring: Lattice):
self._locations = locs


class GeometryObservable(_ElementObservable):
"""Observe the geometrical parameters of the reference trajectory"""
field_list = {'x', 'y', 'angle'}
def __init__(self, refpts: Refpts, param: str,
name: Optional[str] = None,
**kwargs):
# noinspection PyUnresolvedReferences
r"""
Args:
refpts: Observation points.
See ":ref:`Selecting elements in a lattice <refpts>`"
param: Geometry parameter name: one in {'x', 'y', 'angle'}
name: Observable name. If :py:obj:`None`, an explicit
name will be generated
Keyword Args:
statfun: Post-processing function called on the value of the
observable. Example: :pycode:`statfun=numpy.mean`
target: Target value for a constraint. If :py:obj:`None`
(default), the residual will always be zero.
weight: Weight factor: the residual is
:pycode:`((value-target)/weight)**2`
bounds: Tuple of lower and upper bounds. The parameter
is constrained in the interval
[*target*\ +\ *low_bound* *target*\ +\ *up_bound*]
The *target*, *weight* and *bounds* inputs must be broadcastable to the
shape of *value*.
Example:
>>> obs = GeometryObservable(at.Monitor, param='x')
Observe x coordinate of monitors
"""
if param not in self.field_list:
raise ValueError(
f'Expected {param!r} to be one of {self.field_list!r}')
name = self._set_name(name, 'geometry', param)
fun = _recordaccess(param, None)
needs = {Need.GEOMETRY}
super().__init__(fun, refpts, needs=needs, name=name, **kwargs)


class OrbitObservable(_ElementObservable):
"""Observes the transfer matrix at selected locations"""
def __init__(self, refpts: Refpts, axis: AxisDef = None,
Expand Down Expand Up @@ -909,6 +956,8 @@ def obseval(obs):
data.append(trajs[obsrefs[self.passrefs]])
if Need.EMITTANCE in obsneeds:
data.append(emdata)
if Need.GEOMETRY in obsneeds:
data.append(geodata[obsrefs])
obs.evaluate(ring, *data, initial=initial)

@frequency_control
Expand All @@ -918,7 +967,7 @@ def ringeval(ring, dp: Optional[float] = None,
r_in: Orbit = None):
"""Optics computations"""

trajs = orbits = rgdata = eldata = emdata = mxdata = None
trajs = orbits = rgdata = eldata = emdata = mxdata = geodata = None
needs = self.needs

if Need.TRAJECTORY in needs:
Expand Down Expand Up @@ -955,9 +1004,12 @@ def ringeval(ring, dp: Optional[float] = None,
if Need.EMITTANCE in needs:
emdata = ring.envelope_parameters(orbit=o0, keep_lattice=True)

return trajs, orbits, rgdata, eldata, emdata, mxdata
if Need.GEOMETRY in needs:
geodata, _ = ring.get_geometry()

return trajs, orbits, rgdata, eldata, emdata, mxdata, geodata

trajs, orbits, rgdata, eldata, emdata, mxdata = \
trajs, orbits, rgdata, eldata, emdata, mxdata, geodata = \
ringeval(ring, r_in=r_in, **kwargs)

for ob in self:
Expand Down

0 comments on commit a4053cb

Please sign in to comment.