Skip to content

Conversation

@austin-hoover
Copy link
Contributor

This PR adds key word arguments and type hints to all TEAPOT node constructors. Addresses Issue #63.

Changes

  • Added type hints to all arguments in TEAPOT node constructors in orbit.teapot.teapot.
  • Added key word arguments to all TEAPOT node constructors in orbit.teapot.teapot. For example, the kq parameter for QuadTEAPOT is now an argument in the constructor.
  • Added tests to /examples/TEAPOT (where should these go?).

Example: FODO lattice

from orbit.teapot import TEAPOT_Lattice
from orbit.teapot import DriftTEAPOT
from orbit.teapot import QuadTEAPOT

length = 5.0
length_quad = 0.25 * length
length_drift = 0.25 * length
kq = 0.65
nparts = 4

Old:

lattice = TEAPOT_Lattice()

node = DriftTEAPOT("drift1")
node.setLength(length_drift * 0.5)
node.setnParts(nparts)
lattice.addNode(node)

node = QuadTEAPOT("quad1")
node.setLength(length_quad)
node.setnParts(nparts)
node.setParam("kq", +kq)
lattice.addNode(node)

node = DriftTEAPOT("drift2")
node.setLength(length_drift)
node.setnParts(nparts)
lattice.addNode(node)

node = QuadTEAPOT("quad2")
node.setLength(length_quad)
node.setnParts(nparts)
node.setParam("kq", -kq)
lattice.addNode(node)

node = DriftTEAPOT("drift3")
node.setLength(length_drift * 0.5)
node.setnParts(nparts)
lattice.addNode(node)

New:

nodes = [
    DriftTEAPOT(length=length_drift * 0.5, nparts=nparts, name="drift1"),
    QuadTEAPOT(length=length_quad, kq=+kq, nparts=nparts, name="quad1"),
    DriftTEAPOT(length=length_drift, nparts=nparts, name="drift2"),
    QuadTEAPOT(length=length_quad, kq=-kq, nparts=nparts, name="quad2"),
    DriftTEAPOT(length=length_drift * 0.5, nparts=nparts, name="drift3"),
]
lattice = TEAPOT_Lattice()
for node in nodes:
    lattice.addNode(node)

@azukov
Copy link
Member

azukov commented Dec 2, 2025

Good to go.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants