Add dose estimation UI

demo.yaml/beamline.yaml

@@ -48,6 +48,7 @@ objects:
   - sample_view: sample_view.yaml
   - queue_manager: queue_manager.yaml
   - queue_model: queue_model.yaml
+  - dose_estimator: dose_estimator.yaml
 
   # Procedures:
   - beamline_actions: beamline_actions.yaml

demo.yaml/dose_estimator.yaml

@@ -0,0 +1,14 @@
+class: mxcubecore.HardwareObjects.mockup.dose_estimator.MockupDoseEstimator
+configuration:
+    dose_rate_mgy_per_s: 1.0
+    error_probability: 0.1
+    experimental_goals:
+        cryo_highres:
+            label: Cryo High Resolution
+            type: resolution_dependent
+            mgy_per_angstrom: 10
+        room:
+            label: Room Temperature
+            type: static
+            mgy: 0.2
+

mxcubeweb/core/adapter/dose_estimator_adapter.py

@@ -0,0 +1,50 @@
+from typing import ClassVar
+
+from flask import Flask
+from mxcubecore.HardwareObjects.abstract.AbstractDoseEstimator import (
+    AbstractDoseEstimator,
+    DoseEstimateParameters,
+    DoseEstimation,
+)
+
+from mxcubeweb.core.adapter.adapter_base import AdapterBase
+from mxcubeweb.core.models.adaptermodels import HOModel
+from mxcubeweb.core.models.configmodels import ResourceHandlerConfigModel
+
+resource_handler_config = ResourceHandlerConfigModel(
+    commands=["estimate_dose"],
+    attributes=["data", "experimental_goals"],
+)
+
+
+class DoseEstimatorAdapter(AdapterBase):
+    """Adapter for AbstractDoseEstimator hardware objects.
+
+    Exposes dose estimation as a callable command and
+    dose presets as a readable attribute.
+    """
+
+    SUPPORTED_TYPES: ClassVar[list[object]] = [AbstractDoseEstimator]
+
+    def __init__(
+        self,
+        ho: AbstractDoseEstimator,
+        role: str,
+        app: Flask,
+    ) -> None:
+        super().__init__(ho, role, app, resource_handler_config)
+
+    def estimate_dose(
+        self,
+        params: DoseEstimateParameters,
+    ) -> DoseEstimation:
+        return self._ho.estimate_dose(params)
+
+    # Plain `dict` return, type parametrization, i.e.
+    # dict[str, ExperimentalGoal] would not pass as serialization
+    # logic for adapter does not support these.
+    def experimental_goals(self) -> dict:
+        return self._ho.experimental_goals
+
+    def data(self) -> HOModel:
+        return HOModel(**self._dict_repr())

ui/src/actions/taskForm.js

@@ -1,3 +1,5 @@
+import { sendExecuteCommand } from '../api/hardware-object';
+
 export function showTaskForm(
   formName,
   sampleQueueID = -1,
@@ -22,3 +24,26 @@ export function hideTaskParametersForm() {
     type: 'HIDE_FORM',
   };
 }
+
+/**
+ * @typedef {Object} DoseEstimationParameters
+ * @property {number} num_images
+ * @property {number} exp_time_s
+ * @property {number} energy_kev
+ * @property {number} transmission_pct
+ * @property {string} experimental_goal
+ * @property {number} resolution_a
+ *
+ * @param {DoseEstimationParameters} doseParams
+ */
+export function estimateDose(doseParams) {
+  return async (dispatch) => {
+    const result = await sendExecuteCommand(
+      'doseestimator',
+      'dose_estimator',
+      'estimate_dose',
+      doseParams,
+    );
+    dispatch({ type: 'SET_DOSE_ESTIMATE', data: result });
+  };
+}

ui/src/components/Tasks/DataCollection.jsx

@@ -2,11 +2,15 @@
 import React from 'react';
 import { Button, ButtonToolbar, Col, Form, Modal, Row } from 'react-bootstrap';
 import { connect } from 'react-redux';
-import { formValueSelector, reduxForm } from 'redux-form';
+import { change, formValueSelector, reduxForm } from 'redux-form';
 
+import { estimateDose } from '../../actions/taskForm';
 import { SPACE_GROUPS } from '../../constants';
 import { DraggableModal } from '../DraggableModal';
+import TooltipTrigger from '../TooltipTrigger';
 import asyncValidate from './asyncValidate';
+import doseStyles from './DoseEstimation/DoseEstimate.module.css';
+import { computeLimit, DoseLimitInput } from './DoseEstimation/DoseLimitInput';
 import {
   CollapsableRows,
   FieldsHeader,
@@ -278,6 +282,55 @@ class DataCollection extends React.Component {
                 />
               </FieldsRow>
             )}
+            {this.props.canEstimateDose && (
+              <FieldsRow>
+                <Form.Group>
+                  <Row>
+                    <Form.Label column sm={6}>
+                      Estimated dose (MGy)
+                    </Form.Label>
+                    <Form.Label column sm={5}>
+                      {this.props.doseEstimate?.status === 'ok' && (
+                        <span
+                          className={
+                            this.props.doseEstimate.dose_mgy >
+                            this.props.doseLimit
+                              ? doseStyles.over
+                              : doseStyles.under
+                          }
+                        >
+                          {this.props.doseEstimate.dose_mgy?.toFixed(2)}
+                        </span>
+                      )}
+                      {this.props.doseEstimate?.status === 'error' && (
+                        <TooltipTrigger
+                          id="dose-estimate-error-tooltip"
+                          tooltipContent={this.props.doseEstimate.msg}
+                          inModal
+                        >
+                          <span className={doseStyles.unavailable}>N/A</span>
+                        </TooltipTrigger>
+                      )}
+                      {!this.props.doseEstimate && (
+                        <span className={doseStyles.unavailable}>N/A</span>
+                      )}
+                    </Form.Label>
+                  </Row>
+                </Form.Group>
+                <DoseLimitInput
+                  goals={this.props.goals}
+                  resolution={this.props.resolution}
+                  value={this.props.doseLimit || Number.POSITIVE_INFINITY}
+                  currentGoalId={this.props.currentGoalId}
+                  onLimitChange={(limit) => {
+                    this.props.change('dose_limit', limit);
+                  }}
+                  onGoalChange={(goal) =>
+                    this.props.change('exp_goal', goal ?? null)
+                  }
+                />
+              </FieldsRow>
+            )}
             <CollapsableRows>
               <FieldsRow>
                 <InputField propName="kappa" type="number" label="Kappa" />
@@ -352,12 +405,78 @@ const DataCollectionForm = reduxForm({
   validate,
   asyncValidate,
   warn,
+  shouldWarn: ({ values, nextProps, props, initialRender, structure }) => {
+    // Default behaviour + recomputing warnings on estimate changes.
+    if (initialRender) {
+      return true;
+    }
+    if (!structure.deepEqual(values, nextProps.values)) {
+      return true;
+    }
+    return props.doseEstimate !== nextProps.doseEstimate;
+  },
+  onChange: (values, dispatch, props, previousValues) => {
+    // when using dose estimation, changes to form fields
+    // may trigger a need to reestimate.
+    // this handled computations of dependent fields for
+    // dose and dose limit.
+    if (!props.canEstimateDose) {
+      return;
+    }
+    const estimateTriggers = [
+      'energy',
+      'transmission',
+      'num_images',
+      'exp_time',
+      'resolution',
+      'exp_goal',
+      'dose_limit',
+    ];
+    const estimateChanged = estimateTriggers.some(
+      (field) => values[field] !== previousValues[field],
+    );
+    if (estimateChanged) {
+      dispatch(
+        estimateDose({
+          num_images: Number.parseInt(values.num_images, 10),
+          energy_kev: Number.parseFloat(values.energy),
+          transmission_pct: Number.parseFloat(values.transmission),
+          exp_time_s: Number.parseFloat(values.exp_time),
+          resolution_a: Number.parseFloat(values.resolution),
+          experimental_goal: values?.exp_goal,
+          dose_limit_mgy: Number.parseFloat(values?.dose_limit),
+        }),
+      );
+    }
+
+    const currentGoal = props.goals?.[values.exp_goal];
+    if (
+      values.resolution !== previousValues.resolution &&
+      currentGoal?.type === 'resolution_dependent'
+    ) {
+      dispatch(
+        change(
+          'datacollection',
+          'dose_limit',
+          computeLimit(currentGoal, values.resolution),
+        ),
+      );
+    }
+  },
 })(DataCollection);
 
 const selector = formValueSelector('datacollection');
 
 export default connect((state) => {
+  const { doseEstimate } = state.taskForm;
+
+  const doseEstimatorHO = state.beamline.hardwareObjects?.dose_estimator;
   const subdir = selector(state, 'subdir');
+  const canEstimateDose = !!doseEstimatorHO;
+
+  const goals = doseEstimatorHO?.attributes?.experimental_goals || {};
+
+  const resolution = Number.parseFloat(toFixed(state, 'resolution'));
 
   let position = state.taskForm.pointID === '' ? 'PX' : state.taskForm.pointID;
   if (typeof position === 'object') {
@@ -379,6 +498,9 @@ export default connect((state) => {
     state.taskForm.defaultParameters[type.toLowerCase()];
   const { parameters } = state.taskForm.taskData;
 
+  const doseLimit = selector(state, 'dose_limit');
+  const currentGoalId = selector(state, 'exp_goal');
+
   if (Number.parseFloat(parameters.osc_range) === 0) {
     parameters.osc_range =
       state.taskForm.defaultParameters[
@@ -404,11 +526,21 @@ export default connect((state) => {
     filename: fname,
     acqParametersLimits: limits,
     beamline: state.beamline,
+    canEstimateDose,
+    goals,
+    resolution,
+    doseEstimate,
+    doseLimit,
+    currentGoalId,
     detector_mode_list: acq_parameters.detector_mode_list,
     components: state.uiproperties.sample_view_motors.components,
     initialValues: {
       ...parameters,
       beam_size: state.sampleview.currentAperture,
+      dose_limit: Object.values(goals)[0]
+        ? computeLimit(Object.values(goals)[0], resolution)
+        : undefined,
+      exp_goal: Object.keys(goals)[0],
       resolution: toFixed(state, 'resolution'),
       energy: toFixed(state, 'energy'),
       transmission: toFixed(state, 'transmission'),

ui/src/components/Tasks/DoseEstimation/DoseEstimate.module.css

@@ -0,0 +1,15 @@
+.unavailable {
+  text-decoration: underline dashed;
+  cursor: help;
+  color: var(--bs-danger);
+}
+
+.over {
+  color: var(--bs-danger);
+  font-weight: 600;
+}
+
+.under {
+  color: var(--bs-success);
+  font-weight: 600;
+}