diff --git a/src/aind_slims_api/operations/waterlog.py b/src/aind_slims_api/operations/waterlog.py new file mode 100644 index 0000000..762ae44 --- /dev/null +++ b/src/aind_slims_api/operations/waterlog.py @@ -0,0 +1,205 @@ +"""Common operations for logging mouse water and weight with the SLIMS API.""" + +from datetime import datetime +from typing import Optional +import os +import logging + +from pydantic import BaseModel, model_validator, validate_call + +from aind_slims_api import models, SlimsClient, exceptions +from aind_slims_api.models import ( + SlimsMouseContent, + SlimsWaterlogResult, + SlimsWaterRestrictionEvent, + SlimsUser, +) +from aind_slims_api.models.base import SlimsBaseModel + +logger = logging.getLogger(__name__) + + +class WaterlogMouseOperator: + """Class for tracking mouse water/weight data and syncing with SLIMS + + Examples + -------- + + >>> from aind_slims_api import SlimsClient + + >>> client = SlimsClient() + >>> user = client.fetch_model(SlimsUser,username='username') + >>> barcode = '000000' + + >>> mouse = Mouse(barcode,user,client) + + >>> print(f'Mouse {barcode} baseline weight is {mouse.details.baseline_weight_g} g') + + >>> if mouse.details.water_restricted: + ... mouse.switch_to_adlib_water() + + >>> mouse.add_waterlog_record( + ... weight=20, + ... water_earned=0, + ... water_supplement_recommended=1, + ... water_supplement_delivered=1, + ... ) + + >>> mouse.post_baseline_weight(20) + + >>> mouse.switch_to_water_restricted() + + >>> mouse.update_target_weight_fraction(0.9) + + """ + + client: SlimsClient + user: SlimsUser + barcode: str + + details: SlimsMouseContent + waterlog_results: list[SlimsWaterlogResult] + restriction: SlimsWaterRestrictionEvent | None + all_restrictions: list[SlimsWaterRestrictionEvent] + + def __init__(self, barcode: str, user: SlimsUser, slims_client=None): + """Fetch data from Slims for mouse with barcode={mouse_name}""" + self.client = slims_client or SlimsClient() + self.barcode = barcode + self.user = user + + self.details: SlimsMouseContent = None + self.waterlog_results: list[SlimsWaterlogResult] = [] + self.restriction: SlimsWaterRestrictionEvent = None + self.all_restrictions: list[SlimsWaterRestrictionEvent] = [] + + self.link_mouse: str = None + self.link_restrictions: str = None + self.link_wl_records: str = None + + self._fetch_data() + self._fetch_pks() + + def _fetch_data(self): + """Fetches mouse/waterlog/restriction data from SLIMS""" + + self.details = self.client.fetch_model( + SlimsMouseContent, + barcode=self.barcode, + ) + self.waterlog_results = self.client.fetch_models( + SlimsWaterlogResult, + mouse_pk=self.details.pk, + ) + self.all_restrictions = self.client.fetch_models( + SlimsWaterRestrictionEvent, + mouse_pk=self.details.pk, + ) + + if len(self.all_restrictions) > 0: + latest_restriction = self.all_restrictions[-1] + event_active = latest_restriction.end_date is None + if not (self.details.water_restricted == event_active): + logger.warning(f"Warning, inconsistent water restricted data in SLIMS, " f"MID, {self.details.barcode}") + self.restriction = latest_restriction if event_active else None + + self._make_links() + + def _make_links(self): + """Constructs useful links to SLIMS tables""" + self.link_mouse = self.client.rest_link(SlimsMouseContent._slims_table, cntn_cf_labtracksId=self.barcode) + self.link_wl_records = self.client.rest_link(SlimsWaterlogResult._slims_table, rslt_fk_content=self.barcode) + self.link_restrictions = self.client.rest_link( + SlimsWaterRestrictionEvent._slims_table, cnvn_fk_content=self.barcode + ) + + def _fetch_pks(self): + """Fetches useful SLIMS pks""" + self.wrest_pk = self.client.fetch_pk("ContentEventType", cnvt_uniqueIdentifier="cnvt_water_restriction") + self.wl_test_pk = self.client.fetch_pk("Test", test_name="test_waterlog") + + def add_waterlog_record( + self, + weight: float, + water_earned: float, + water_supplement_recommended: float, + water_supplement_delivered: float, + comments: Optional[str] = None, + workstation: Optional[str] = None, + ): + """Creates and adds a new waterlog weight/water record to SLIMS, and + updates self.waterlog_results accordingly""" + + total_water = water_earned + water_supplement_delivered + + record = SlimsWaterlogResult( + weight_g=weight, + water_earned_ml=water_earned, + water_supplement_recommended_ml=water_supplement_recommended, + water_supplement_delivered_ml=water_supplement_delivered, + total_water_ml=total_water, + comments=comments, + test_pk=self.wl_test_pk, + mouse_pk=self.details.pk, + workstation=workstation or os.environ.get("aibs_comp_id", None), + ) + + record = self.client.add_model(record) + self.waterlog_results.append(record) + logger.info("Added SLIMS Waterlog record") + + def post_baseline_weight(self, new_baseline_weight: float): + """Update the baseline weight in SLIMS and self.details""" + self.details.baseline_weight_g = new_baseline_weight + self.details = self.client.update_model(self.details, "baseline_weight_g") + logger.info(f"Updated mouse {self.barcode} " f"baseline weight to {new_baseline_weight}") + + def switch_to_water_restricted(self, target_weight_fraction: float): + """If the mouse is on ad-lib water, + - Create a water restriction event starting today + - Update the mouse's water_restricted field + - Update SLIMS with the above + - Update local data accordingly""" + if self.details.water_restricted: + logger.info("Mouse is already water restricted") + return + + new_restriction = SlimsWaterRestrictionEvent( + assigned_by=self.user.full_name, + target_weight_fraction=target_weight_fraction, + mouse_pk=self.details.pk, + cnvn_fk_contentEventType=self.wrest_pk, + ) + self.details.water_restricted = True + self.details = self.client.update_model(self.details, "water_restricted") + + self.restriction = self.client.add_model(new_restriction) + self.all_restrictions.append(self.restriction) + logger.info(f"Switched mouse {self.barcode} to Water Restricted") + + def switch_to_adlib_water(self): + """If the mouse is water restricted, + - Set the end date of the active restriction event to today + - Update the mouse's water_restricted field + - Update SLIMS with the above + - Update local data accordingly""" + if not self.details.water_restricted: + logger.info("Mouse is already on ad-lib water") + return + + self.details.water_restricted = False + self.details = self.client.update_model(self.details, "water_restricted") + + self.restriction.end_date = datetime.now() + self.restriction = self.client.update_model(self.restriction, "end_date") + self.restriction = None + logger.info(f"Switched mouse {self.barcode} to Adlib Water") + + def update_target_weight_fraction(self, new_twf: float): + """Update the target weight fraction of the active restriction""" + if not self.details.water_restricted: + logger.info("Mouse is not water restricted") + return + self.restriction.target_weight_fraction = new_twf + self.restriction = self.client.update_model(self.restriction, "target_weight_fraction") + logger.info(f"Updated mouse {self.barcode} " f"target weight fraction to {new_twf}") diff --git a/tests/test_waterlog.py b/tests/test_waterlog.py new file mode 100644 index 0000000..e69de29