"""
Object definition for Source, Beam and Attenuator
"""
from __future__ import annotations
from collections import namedtuple
from dataclasses import dataclass
from dataclasses_json import DataClassJsonMixin
# Describe facility
Facility = namedtuple("Facility", ("name", "short_name", "type", "id"))
Facility.__doc__ = """Facility description"""
[docs]
class Source:
"""Source definition."""
def __init__(
self,
beamline: str,
facility: Facility = Facility(
"Diamond Light Source",
"DLS",
"Synchrotron X-ray Source",
None,
),
probe: str | None = None,
):
self.beamline = beamline
self.name = facility.name
self.short_name = facility.short_name
self.facility_type = facility.type
self.facility_id = facility.id
self.probe = probe
def __repr__(self) -> str:
msg = f"Facility: {self.name} - {self.facility_type}. \n\t"
msg += f"Beamline / instrument: {self.beamline} \n\t"
if self.probe:
msg += "Probe: {}"
return f"Source information: \n\t{msg}"
[docs]
def set_instrument_name(self) -> str:
"""Set the instrument name from the details saved in source.
If source type is not defined a priori, the function will assume it is a Synchrotron.
If the facility_id is defined inside the source dictionary, that is the value that will be used.
Naming tries to follow the recommended convention for NXmx:
https://mmcif.wwpdb.org/dictionaries/mmcif_pdbx_v50.dic/Items/_diffrn_source.type.html
Returns:
name (str): The name to write under '/entry/instrument/name'
"""
facility_id = "DIAMOND" if self.facility_id is None else self.facility_id
if (
self.facility_type is None
or "SYNCHROTRON" not in self.facility_type.upper()
):
return f"{facility_id} {self.beamline}"
else:
return f"{facility_id} BEAMLINE {self.beamline}"
def _generate_source_dict(self):
source = {
"name": self.name,
"type": self.facility_type,
"short_name": self.short_name,
"beamline_name": self.beamline,
}
if self.probe:
source["probe"] = self.probe
if self.facility_id:
source["facility_id"] = self.facility_id
return source
[docs]
def to_dict(self):
"""Write source information to a dictionary."""
return self._generate_source_dict()
[docs]
@dataclass
class Beam(DataClassJsonMixin):
"""Beam definition."""
wavelength: float
flux: float | None = None
[docs]
@dataclass
class Attenuator(DataClassJsonMixin):
"""Attenuator definition."""
transmission: float