Skip to content

Commit

Permalink
Updaed pydantic
Browse files Browse the repository at this point in the history
  • Loading branch information
jmp1985 committed Jun 10, 2024
1 parent 6e53564 commit b31cf7b
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 22 deletions.
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ def main():
"pandas",
"pillow",
"profet",
"pydantic>=1.10.13",
"pydantic",
"python-multem",
"pyyaml",
"scipy",
Expand Down
43 changes: 22 additions & 21 deletions src/parakeet/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
import yaml

from enum import Enum
from pydantic import BaseModel as PydanticBaseModel
from pydantic import ConfigDict, BaseModel as PydanticBaseModel
from pydantic import Field
from typing import List
from typing import Optional
Expand All @@ -40,12 +40,7 @@ class BaseModel(PydanticBaseModel):
"""

class Config:
# Ensure that enums use string values
use_enum_values = True

# Don't allow extra fields
extra = "forbid"
model_config = ConfigDict(use_enum_values=True, extra="forbid")


class Auto(str, Enum):
Expand Down Expand Up @@ -104,11 +99,11 @@ class Cylinder(BaseModel):
(0, 1, 0), description="The axis of the cylinder"
)

offset_x: List[float] = Field(
offset_x: Optional[List[float]] = Field(
None, description="The x offset as a function of cylinder y position"
)

offset_z: List[float] = Field(
offset_z: Optional[List[float]] = Field(
None, description="The z offset as a function of cylinder y position"
)

Expand Down Expand Up @@ -157,6 +152,7 @@ class MoleculePose(BaseModel):
"""

position: Optional[Tuple[float, float, float]] = Field(
None,
description=(
"The molecule position (A, A, A). Setting this to null or an "
"empty list will cause parakeet to give a random position. "
Expand All @@ -170,6 +166,7 @@ class MoleculePose(BaseModel):
)

orientation: Optional[Tuple[float, float, float]] = Field(
None,
description=(
"The molecule orientation defined as a rotation vector where "
"the direction of the vector gives the rotation axis and the "
Expand All @@ -191,7 +188,7 @@ class CoordinateFile(BaseModel):
"""

filename: str = Field(
filename: Optional[str] = Field(
None, description="The filename of the atomic coordinates to use (*.pdb, *.cif)"
)

Expand All @@ -200,6 +197,7 @@ class CoordinateFile(BaseModel):
scale: float = Field(1, description="Scale the coordinates x' = x * scale")

position: Optional[Tuple[float, float, float]] = Field(
None,
description=(
"The model position (A, A, A). "
"If recentre if set then the model will be centred on the given position. "
Expand All @@ -212,6 +210,7 @@ class CoordinateFile(BaseModel):
)

orientation: Optional[Tuple[float, float, float]] = Field(
None,
description=(
"The model orientation defined as a rotation vector where "
"the direction of the vector gives the rotation axis and the "
Expand Down Expand Up @@ -292,11 +291,11 @@ class Molecules(BaseModel):
"""

local: Optional[List[LocalMolecule]] = Field(
description="The local molecules to include in the sample model"
None, description="The local molecules to include in the sample model"
)

pdb: Optional[List[PDBMolecule]] = Field(
description="The PDB molecules to include in the sample model"
None, description="The PDB molecules to include in the sample model"
)


Expand Down Expand Up @@ -346,17 +345,17 @@ class Sample(BaseModel):
)

coords: Optional[CoordinateFile] = Field(
description="Coordinates to initialise the sample"
None, description="Coordinates to initialise the sample"
)

molecules: Optional[Molecules] = Field(
description="The molecules to include in the sample model"
None, description="The molecules to include in the sample model"
)

ice: Optional[Ice] = Field(description="The atomic ice model parameters.")
ice: Optional[Ice] = Field(None, description="The atomic ice model parameters.")

sputter: Optional[Sputter] = Field(
description="The sputter coating model parameters."
None, description="The sputter coating model parameters."
)


Expand Down Expand Up @@ -500,7 +499,7 @@ class Microscope(BaseModel):
"""

model: MicroscopeModel = Field(
model: Optional[MicroscopeModel] = Field(
None, description="Use parameters for a given microscope model"
)

Expand All @@ -512,7 +511,7 @@ class Microscope(BaseModel):
PhasePlate(), description="The phase plate parameters"
)

objective_aperture_cutoff_freq: float = Field(
objective_aperture_cutoff_freq: Optional[float] = Field(
None, description="The objective aperture cutoff frequency (1/A)"
)

Expand Down Expand Up @@ -643,7 +642,7 @@ class Scan(BaseModel):
),
)

drift: Optional[Drift] = Field(description="The drift model parameters")
drift: Optional[Drift] = Field(None, description="The drift model parameters")


class InelasticModel(str, Enum):
Expand Down Expand Up @@ -709,11 +708,13 @@ class Simulation(BaseModel):
False, description="Use the radiation damage model (True/False)"
)

inelastic_model: InelasticModel = Field(
inelastic_model: Optional[InelasticModel] = Field(
None, description="The inelastic model parameters"
)

mp_loss_width: float = Field(None, description="The MPL energy filter width")
mp_loss_width: Optional[float] = Field(
None, description="The MPL energy filter width"
)

mp_loss_position: MPLPosition = Field(
"peak", description="The MPL energy filter position"
Expand Down

0 comments on commit b31cf7b

Please sign in to comment.