qmbed/ad/
mod.rs

1//! Rust-native differentiation primitives for QMBED scientific operations.
2//!
3//! The native layer owns the mathematical derivative semantics. Optional
4//! adapters such as the optional `chainrules` module only translate these tested primitives into
5//! another protocol; they do not reimplement the derivatives.
6
7use std::collections::HashSet;
8use std::sync::Arc;
9
10use crate::operator::{LinearOperator, MatrixFormat, Operator, QuantumOperator};
11use crate::runtime::{Runtime, RuntimeAdjointLinearOperator, RuntimeBuffer, RuntimeLinearOperator};
12use crate::solve::{EigshOptions, EigshWorkspace, SpectrumTarget, eigsh_with_workspace};
13use crate::{Complex64, QmbedError, Result};
14
15/// Differential domain of one ordered operator parameter.
16#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17pub enum ParameterDomain {
18    /// A real parameter represented by a complex value with zero imaginary part.
19    Real,
20    /// A complex parameter differentiated under the real pairing `Re(x†y)`.
21    Complex,
22}
23
24/// Stable names and differential domains for an ordered parameter vector.
25#[derive(Clone, Debug, Eq, PartialEq)]
26pub struct ParameterSchema {
27    names: Arc<[String]>,
28    domains: Arc<[ParameterDomain]>,
29}
30
31impl ParameterSchema {
32    /// Construct a schema after validating names and domain count.
33    pub fn new(
34        names: impl IntoIterator<Item = String>,
35        domains: impl IntoIterator<Item = ParameterDomain>,
36    ) -> Result<Self> {
37        let names: Vec<_> = names.into_iter().collect();
38        let domains: Vec<_> = domains.into_iter().collect();
39        if names.len() != domains.len() {
40            return Err(QmbedError::DimensionMismatch(format!(
41                "parameter schema has {} names and {} domains",
42                names.len(),
43                domains.len()
44            )));
45        }
46        let mut unique = HashSet::with_capacity(names.len());
47        if let Some(name) = names
48            .iter()
49            .find(|name| name.is_empty() || !unique.insert(name.as_str()))
50        {
51            return Err(QmbedError::InvalidOptions(format!(
52                "parameter names must be nonempty and unique; invalid name {name:?}"
53            )));
54        }
55        Ok(Self {
56            names: names.into(),
57            domains: domains.into(),
58        })
59    }
60
61    /// Create a uniform schema from a [`QuantumOperator`]'s stable component order.
62    pub fn for_operator(operator: &QuantumOperator, domain: ParameterDomain) -> Self {
63        let names: Vec<_> = operator.component_names().map(str::to_owned).collect();
64        let domains = vec![domain; names.len()];
65        Self {
66            names: names.into(),
67            domains: domains.into(),
68        }
69    }
70
71    /// Create a mixed real/complex schema for an operator.
72    pub fn for_operator_with_domains(
73        operator: &QuantumOperator,
74        domains: impl IntoIterator<Item = ParameterDomain>,
75    ) -> Result<Self> {
76        Self::new(operator.component_names().map(str::to_owned), domains)
77    }
78
79    /// Ordered parameter names.
80    pub fn names(&self) -> &[String] {
81        &self.names
82    }
83
84    /// Ordered differential domains.
85    pub fn domains(&self) -> &[ParameterDomain] {
86        &self.domains
87    }
88
89    /// Number of parameters.
90    pub fn len(&self) -> usize {
91        self.names.len()
92    }
93
94    /// Whether the schema is empty.
95    pub fn is_empty(&self) -> bool {
96        self.names.is_empty()
97    }
98
99    /// Return the stable index of one parameter name.
100    pub fn index_of(&self, name: &str) -> Option<usize> {
101        self.names.iter().position(|candidate| candidate == name)
102    }
103}
104
105/// Ordered primal values for a parameterized operator.
106#[derive(Clone, Debug, PartialEq)]
107pub struct ParameterValues {
108    schema: Arc<ParameterSchema>,
109    values: Vec<Complex64>,
110}
111
112impl ParameterValues {
113    /// Construct values under an existing schema.
114    pub fn new(
115        schema: Arc<ParameterSchema>,
116        values: impl IntoIterator<Item = Complex64>,
117    ) -> Result<Self> {
118        let values: Vec<_> = values.into_iter().collect();
119        validate_parameter_vector(&schema, &values, "parameter values")?;
120        Ok(Self { schema, values })
121    }
122
123    /// Construct an all-real parameter vector in operator component order.
124    pub fn real(operator: &QuantumOperator, values: impl IntoIterator<Item = f64>) -> Result<Self> {
125        let schema = Arc::new(ParameterSchema::for_operator(
126            operator,
127            ParameterDomain::Real,
128        ));
129        Self::new(
130            schema,
131            values.into_iter().map(|value| Complex64::new(value, 0.0)),
132        )
133    }
134
135    /// Construct an all-complex parameter vector in operator component order.
136    pub fn complex(
137        operator: &QuantumOperator,
138        values: impl IntoIterator<Item = Complex64>,
139    ) -> Result<Self> {
140        let schema = Arc::new(ParameterSchema::for_operator(
141            operator,
142            ParameterDomain::Complex,
143        ));
144        Self::new(schema, values)
145    }
146
147    /// Shared parameter schema.
148    pub fn schema(&self) -> &Arc<ParameterSchema> {
149        &self.schema
150    }
151
152    /// Ordered numeric values.
153    pub fn values(&self) -> &[Complex64] {
154        &self.values
155    }
156
157    /// Construct a direction in the same differential space.
158    pub fn direction(
159        &self,
160        values: impl IntoIterator<Item = Complex64>,
161    ) -> Result<ParameterDirection> {
162        ParameterDirection::new(Arc::clone(&self.schema), values)
163    }
164
165    /// Return one named value.
166    pub fn get(&self, name: &str) -> Option<Complex64> {
167        self.schema.index_of(name).map(|index| self.values[index])
168    }
169}
170
171/// Ordered forward perturbation of [`ParameterValues`].
172#[derive(Clone, Debug, PartialEq)]
173pub struct ParameterDirection {
174    schema: Arc<ParameterSchema>,
175    values: Vec<Complex64>,
176}
177
178impl ParameterDirection {
179    /// Construct a direction after checking its schema and domains.
180    pub fn new(
181        schema: Arc<ParameterSchema>,
182        values: impl IntoIterator<Item = Complex64>,
183    ) -> Result<Self> {
184        let values: Vec<_> = values.into_iter().collect();
185        validate_parameter_vector(&schema, &values, "parameter direction")?;
186        Ok(Self { schema, values })
187    }
188
189    /// Shared parameter schema.
190    pub fn schema(&self) -> &Arc<ParameterSchema> {
191        &self.schema
192    }
193
194    /// Ordered direction values.
195    pub fn values(&self) -> &[Complex64] {
196        &self.values
197    }
198}
199
200/// Ordered reverse sensitivity of [`ParameterValues`].
201#[derive(Clone, Debug, PartialEq)]
202pub struct ParameterGradient {
203    schema: Arc<ParameterSchema>,
204    values: Vec<Complex64>,
205}
206
207impl ParameterGradient {
208    fn new(schema: Arc<ParameterSchema>, values: Vec<Complex64>) -> Self {
209        Self { schema, values }
210    }
211
212    /// Shared parameter schema.
213    pub fn schema(&self) -> &Arc<ParameterSchema> {
214        &self.schema
215    }
216
217    /// Ordered gradient values.
218    pub fn values(&self) -> &[Complex64] {
219        &self.values
220    }
221
222    /// Return one named gradient.
223    pub fn get(&self, name: &str) -> Option<Complex64> {
224        self.schema.index_of(name).map(|index| self.values[index])
225    }
226}
227
228/// Reliability classification attached to a scientific gradient.
229#[derive(Clone, Copy, Debug, Eq, PartialEq)]
230pub enum GradientStatus {
231    /// The rule is exact or all requested residual and conditioning checks passed.
232    Reliable,
233    /// The derivative exists, but diagnostics indicate poor conditioning.
234    IllConditioned,
235    /// The requested mathematical derivative is undefined.
236    NonDifferentiable,
237}
238
239/// Backend-independent evidence returned with a native derivative.
240#[derive(Clone, Debug, PartialEq)]
241pub struct GradientDiagnostics {
242    /// Reliability classification.
243    pub status: GradientStatus,
244    /// Optional primal numerical residual.
245    pub primal_residual: Option<f64>,
246    /// Optional reverse numerical residual.
247    pub backward_residual: Option<f64>,
248    /// Optional spectral gap or separation diagnostic.
249    pub spectral_gap: Option<f64>,
250    /// Number of component actions used by the primal/JVP computation.
251    pub primal_applications: usize,
252    /// Number of component actions used by the reverse computation.
253    pub backward_applications: usize,
254    /// Number of forward states or actions recomputed by the reverse pass.
255    pub recomputations: usize,
256}
257
258/// Ground-state energy, Hellmann--Feynman gradient, and solver evidence.
259///
260/// This rule differentiates an isolated lowest eigenvalue of a real-parameter
261/// Hermitian operator family. It deliberately does not return a derivative of
262/// an arbitrarily phased eigenvector.
263#[derive(Clone, Debug)]
264pub struct GroundStateEnergyGradient {
265    /// Algebraically smallest eigenvalue.
266    pub energy: f64,
267    /// Normalized eigenvector used by the Hellmann--Feynman rule.
268    pub state: Vec<Complex64>,
269    /// Ordered `dE/dθᵢ = ⟨ψ|Aᵢ|ψ⟩`.
270    pub gradient: ParameterGradient,
271    /// Residual and spectral-separation evidence.
272    pub diagnostics: GradientDiagnostics,
273    /// Eigensolver iteration count.
274    pub eigensolver_iterations: usize,
275}
276
277struct BoundQuantumOperator<'a> {
278    family: &'a QuantumOperator,
279    coefficients: &'a [Complex64],
280}
281
282impl LinearOperator for BoundQuantumOperator<'_> {
283    fn shape(&self) -> (usize, usize) {
284        self.family.shape()
285    }
286
287    fn format(&self) -> MatrixFormat {
288        MatrixFormat::MatrixFree
289    }
290
291    fn apply(&self, input: &[Complex64], output: &mut [Complex64]) -> Result<()> {
292        self.family
293            .apply_coefficients(self.coefficients, input, output)
294    }
295
296    fn is_real(&self) -> bool {
297        self.coefficients.iter().all(|value| value.im == 0.0)
298            && self
299                .family
300                .components()
301                .iter()
302                .all(|component| component.operator().is_real())
303    }
304}
305
306/// Differentiate the isolated algebraic ground-state energy.
307///
308/// The forward solve requests at least two eigenpairs so the returned
309/// diagnostics can expose the spectral gap. Reverse work then requires one
310/// component action per parameter, instead of the two additional eigensolves
311/// per parameter required by central finite differences.
312///
313/// All parameters and every component must be real and Hermitian. Degenerate
314/// or numerically unresolved ground states are reported as
315/// [`GradientStatus::NonDifferentiable`] rather than silently assigning a
316/// derivative to an arbitrary eigenvector.
317pub fn ground_state_energy_gradient(
318    operator: &QuantumOperator,
319    parameters: &ParameterValues,
320    options: EigshOptions,
321    workspace: &mut EigshWorkspace,
322) -> Result<GroundStateEnergyGradient> {
323    validate_ground_state_arguments(operator, parameters, &options)?;
324    let bound = BoundQuantumOperator {
325        family: operator,
326        coefficients: parameters.values(),
327    };
328    let eigensystem = eigsh_with_workspace(&bound, options, workspace)?;
329    let energy = eigensystem.eigenvalues[0];
330    let gap = eigensystem.eigenvalues[1] - energy;
331    let state = eigensystem.eigenvectors[0].clone();
332    let mut contribution = vec![Complex64::new(0.0, 0.0); state.len()];
333    let mut gradient = Vec::with_capacity(operator.components().len());
334    for component in operator.components() {
335        component.operator().apply(&state, &mut contribution)?;
336        let derivative = state
337            .iter()
338            .zip(&contribution)
339            .map(|(left, right)| left.conj() * right)
340            .sum::<Complex64>();
341        gradient.push(Complex64::new(derivative.re, 0.0));
342    }
343
344    let residual = eigensystem.residuals.iter().copied().fold(0.0, f64::max);
345    let scale = energy.abs().max(1.0);
346    let unresolved = (10.0 * residual).max(64.0 * f64::EPSILON * scale);
347    let status = if gap <= unresolved {
348        GradientStatus::NonDifferentiable
349    } else if gap <= scale * 1.0e-8 {
350        GradientStatus::IllConditioned
351    } else {
352        GradientStatus::Reliable
353    };
354    Ok(GroundStateEnergyGradient {
355        energy,
356        state,
357        gradient: ParameterGradient::new(Arc::clone(parameters.schema()), gradient),
358        diagnostics: GradientDiagnostics {
359            status,
360            primal_residual: Some(residual),
361            backward_residual: None,
362            spectral_gap: Some(gap),
363            primal_applications: 0,
364            backward_applications: operator.components().len(),
365            recomputations: 0,
366        },
367        eigensolver_iterations: eigensystem.iterations,
368    })
369}
370
371fn validate_ground_state_arguments(
372    operator: &QuantumOperator,
373    parameters: &ParameterValues,
374    options: &EigshOptions,
375) -> Result<()> {
376    validate_operator_schema(operator, parameters)?;
377    let (rows, columns) = operator.shape();
378    if rows != columns {
379        return Err(QmbedError::DimensionMismatch(
380            "ground-state differentiation requires a square operator family".into(),
381        ));
382    }
383    if options.eigenpairs < 2 {
384        return Err(QmbedError::InvalidOptions(
385            "ground-state differentiation requires at least two eigenpairs for gap diagnostics"
386                .into(),
387        ));
388    }
389    if !matches!(options.target, SpectrumTarget::SmallestAlgebraic) {
390        return Err(QmbedError::InvalidOptions(
391            "ground-state differentiation requires the SmallestAlgebraic target".into(),
392        ));
393    }
394    if parameters
395        .schema()
396        .domains()
397        .iter()
398        .any(|domain| *domain != ParameterDomain::Real)
399    {
400        return Err(QmbedError::InvalidOptions(
401            "ground-state energy gradients currently require real parameters".into(),
402        ));
403    }
404    if let Some(component) = operator
405        .components()
406        .iter()
407        .find(|component| !component.operator().is_hermitian(1.0e-12))
408    {
409        return Err(QmbedError::InvalidOptions(format!(
410            "operator component {:?} is not Hermitian",
411            component.name()
412        )));
413    }
414    Ok(())
415}
416
417impl GradientDiagnostics {
418    fn exact_operator(primal_applications: usize, backward_applications: usize) -> Self {
419        Self {
420            status: GradientStatus::Reliable,
421            primal_residual: Some(0.0),
422            backward_residual: Some(0.0),
423            spectral_gap: None,
424            primal_applications,
425            backward_applications,
426            recomputations: backward_applications / 2,
427        }
428    }
429}
430
431/// Primal value and forward tangent from parameterized operator application.
432#[derive(Clone, Debug)]
433pub struct ApplyJvp<B> {
434    /// `A(θ)x`.
435    pub value: B,
436    /// `A(θ)dx + dA(θ)[dθ]x`.
437    pub tangent: B,
438    /// Work and reliability evidence.
439    pub diagnostics: GradientDiagnostics,
440}
441
442/// State and parameter cotangents returned by an operator pullback.
443#[derive(Clone, Debug)]
444pub struct ApplyCotangents<B> {
445    /// Cotangent with respect to the ordered parameters.
446    pub parameters: ParameterGradient,
447    /// Cotangent with respect to the input state.
448    pub state: B,
449    /// Work and reliability evidence.
450    pub diagnostics: GradientDiagnostics,
451}
452
453/// Evaluate a parameterized operator and its JVP without materializing `A(θ)`.
454pub fn apply_jvp<R>(
455    runtime: &R,
456    operator: &QuantumOperator,
457    parameters: &ParameterValues,
458    parameter_direction: &ParameterDirection,
459    input: &R::Buffer,
460    input_direction: &R::Buffer,
461) -> Result<ApplyJvp<R::Buffer>>
462where
463    R: Runtime,
464    Operator: RuntimeLinearOperator<R>,
465{
466    validate_operator_arguments::<R>(
467        operator,
468        parameters,
469        Some(parameter_direction),
470        input,
471        Some(input_direction),
472    )?;
473
474    let (rows, _) = operator.shape();
475    let mut value = runtime.zeros(rows)?;
476    let mut tangent = runtime.zeros(rows)?;
477    let mut input_contribution = runtime.zeros(rows)?;
478    let mut direction_contribution = runtime.zeros(rows)?;
479
480    for ((component, coefficient), parameter_tangent) in operator
481        .components()
482        .iter()
483        .zip(parameters.values())
484        .zip(parameter_direction.values())
485    {
486        runtime.fill(&mut input_contribution, Complex64::new(0.0, 0.0))?;
487        component
488            .operator()
489            .apply_on(runtime, input, &mut input_contribution)?;
490        runtime.axpy(*coefficient, &input_contribution, &mut value)?;
491        runtime.axpy(*parameter_tangent, &input_contribution, &mut tangent)?;
492
493        runtime.fill(&mut direction_contribution, Complex64::new(0.0, 0.0))?;
494        component
495            .operator()
496            .apply_on(runtime, input_direction, &mut direction_contribution)?;
497        runtime.axpy(*coefficient, &direction_contribution, &mut tangent)?;
498    }
499
500    Ok(ApplyJvp {
501        value,
502        tangent,
503        diagnostics: GradientDiagnostics::exact_operator(2 * operator.components().len(), 0),
504    })
505}
506
507/// Evaluate a parameterized operator and prepare a one-shot reverse pullback.
508pub fn apply_vjp<'a, R>(
509    runtime: &'a R,
510    operator: &'a QuantumOperator,
511    parameters: &'a ParameterValues,
512    input: &'a R::Buffer,
513) -> Result<(R::Buffer, ApplyPullback<'a, R>)>
514where
515    R: Runtime,
516    Operator: RuntimeLinearOperator<R> + RuntimeAdjointLinearOperator<R>,
517{
518    validate_operator_arguments::<R>(operator, parameters, None, input, None)?;
519    let (rows, _) = operator.shape();
520    let mut value = runtime.zeros(rows)?;
521    let mut contribution = runtime.zeros(rows)?;
522    for (component, coefficient) in operator.components().iter().zip(parameters.values()) {
523        runtime.fill(&mut contribution, Complex64::new(0.0, 0.0))?;
524        component
525            .operator()
526            .apply_on(runtime, input, &mut contribution)?;
527        runtime.axpy(*coefficient, &contribution, &mut value)?;
528    }
529    let pullback = ApplyPullback {
530        runtime,
531        operator,
532        parameters,
533        input,
534    };
535    Ok((value, pullback))
536}
537
538/// One-shot pullback for [`apply_vjp`].
539pub struct ApplyPullback<'a, R>
540where
541    R: Runtime,
542{
543    runtime: &'a R,
544    operator: &'a QuantumOperator,
545    parameters: &'a ParameterValues,
546    input: &'a R::Buffer,
547}
548
549impl<R> ApplyPullback<'_, R>
550where
551    R: Runtime,
552    Operator: RuntimeLinearOperator<R> + RuntimeAdjointLinearOperator<R>,
553{
554    /// Pull an output cotangent back to state and parameter cotangents.
555    pub fn backward(self, output_cotangent: &R::Buffer) -> Result<ApplyCotangents<R::Buffer>> {
556        let (rows, columns) = self.operator.shape();
557        if output_cotangent.len() != rows {
558            return Err(QmbedError::DimensionMismatch(format!(
559                "operator pullback requires output cotangent length {rows}, got {}",
560                output_cotangent.len()
561            )));
562        }
563
564        let mut state = self.runtime.zeros(columns)?;
565        let mut state_contribution = self.runtime.zeros(columns)?;
566        let mut parameter_contribution = self.runtime.zeros(rows)?;
567        let mut gradient = Vec::with_capacity(self.operator.components().len());
568
569        for ((component, coefficient), domain) in self
570            .operator
571            .components()
572            .iter()
573            .zip(self.parameters.values())
574            .zip(self.parameters.schema.domains())
575        {
576            self.runtime
577                .fill(&mut state_contribution, Complex64::new(0.0, 0.0))?;
578            component.operator().apply_adjoint_on(
579                self.runtime,
580                output_cotangent,
581                &mut state_contribution,
582            )?;
583            self.runtime
584                .axpy(coefficient.conj(), &state_contribution, &mut state)?;
585
586            self.runtime
587                .fill(&mut parameter_contribution, Complex64::new(0.0, 0.0))?;
588            component
589                .operator()
590                .apply_on(self.runtime, self.input, &mut parameter_contribution)?;
591            let cotangent = self
592                .runtime
593                .dotc(&parameter_contribution, output_cotangent)?;
594            gradient.push(match domain {
595                ParameterDomain::Real => Complex64::new(cotangent.re, 0.0),
596                ParameterDomain::Complex => cotangent,
597            });
598        }
599
600        let applications = 2 * self.operator.components().len();
601        Ok(ApplyCotangents {
602            parameters: ParameterGradient::new(Arc::clone(&self.parameters.schema), gradient),
603            state,
604            diagnostics: GradientDiagnostics::exact_operator(
605                self.operator.components().len(),
606                applications,
607            ),
608        })
609    }
610}
611
612fn validate_parameter_vector(
613    schema: &ParameterSchema,
614    values: &[Complex64],
615    label: &str,
616) -> Result<()> {
617    if values.len() != schema.len() {
618        return Err(QmbedError::DimensionMismatch(format!(
619            "{label} has length {}, expected {}",
620            values.len(),
621            schema.len()
622        )));
623    }
624    for ((name, domain), value) in schema.names().iter().zip(schema.domains()).zip(values) {
625        if !value.re.is_finite() || !value.im.is_finite() {
626            return Err(QmbedError::InvalidOptions(format!(
627                "{label} for {name:?} must be finite"
628            )));
629        }
630        if matches!(domain, ParameterDomain::Real) && value.im != 0.0 {
631            return Err(QmbedError::InvalidOptions(format!(
632                "{label} for real parameter {name:?} must have zero imaginary part"
633            )));
634        }
635    }
636    Ok(())
637}
638
639fn validate_operator_arguments<R>(
640    operator: &QuantumOperator,
641    parameters: &ParameterValues,
642    direction: Option<&ParameterDirection>,
643    input: &R::Buffer,
644    input_direction: Option<&R::Buffer>,
645) -> Result<()>
646where
647    R: Runtime,
648{
649    validate_operator_schema(operator, parameters)?;
650    if let Some(direction) = direction {
651        if direction.schema.as_ref() != parameters.schema.as_ref() {
652            return Err(QmbedError::InvalidOptions(
653                "parameter direction uses a different schema".into(),
654            ));
655        }
656    }
657    let (_, columns) = operator.shape();
658    if input.len() != columns {
659        return Err(QmbedError::DimensionMismatch(format!(
660            "parameterized operator requires input length {columns}, got {}",
661            input.len()
662        )));
663    }
664    if let Some(input_direction) = input_direction {
665        if input_direction.len() != columns {
666            return Err(QmbedError::DimensionMismatch(format!(
667                "operator JVP requires input direction length {columns}, got {}",
668                input_direction.len()
669            )));
670        }
671    }
672    Ok(())
673}
674
675fn validate_operator_schema(
676    operator: &QuantumOperator,
677    parameters: &ParameterValues,
678) -> Result<()> {
679    let operator_names: Vec<_> = operator.component_names().collect();
680    if operator_names.len() != parameters.schema.len()
681        || operator_names
682            .iter()
683            .zip(parameters.schema.names())
684            .any(|(operator_name, parameter_name)| operator_name != parameter_name)
685    {
686        return Err(QmbedError::InvalidOptions(
687            "parameter schema does not match QuantumOperator component order".into(),
688        ));
689    }
690    Ok(())
691}
692
693/// `chainrules-core 0.2` adapters for the Rust-native operator rules.
694#[cfg(feature = "chainrules")]
695pub mod chainrules {
696    use std::sync::Arc;
697
698    use chainrules_core::{Differentiable, JvpRule, Pullback, VjpRule};
699
700    use super::{
701        ApplyCotangents, ApplyPullback, GradientStatus, ParameterDirection, ParameterGradient,
702        ParameterValues, apply_jvp, apply_vjp, ground_state_energy_gradient,
703    };
704    use crate::operator::{Operator, QuantumOperator};
705    use crate::runtime::{Runtime, RuntimeAdjointLinearOperator, RuntimeLinearOperator};
706    use crate::solve::{EigshOptions, EigshWorkspace};
707    use crate::{QmbedError, Result};
708
709    impl Differentiable for ParameterValues {
710        type Tangent = ParameterDirection;
711        type Cotangent = ParameterGradient;
712    }
713
714    /// Runtime-owned state returned through `chainrules-core`.
715    #[derive(Clone, Debug)]
716    pub struct State<B>(pub B);
717
718    impl<B> State<B> {
719        /// Consume the wrapper.
720        pub fn into_inner(self) -> B {
721            self.0
722        }
723    }
724
725    impl<B> Differentiable for State<B> {
726        type Tangent = B;
727        type Cotangent = B;
728    }
729
730    /// Borrowed primal arguments for parameterized operator application.
731    pub struct ApplyArguments<'a, R>
732    where
733        R: Runtime,
734    {
735        /// Ordered operator parameters.
736        pub parameters: &'a ParameterValues,
737        /// Runtime-owned input state.
738        pub state: &'a R::Buffer,
739    }
740
741    /// Borrowed forward perturbations corresponding to [`ApplyArguments`].
742    pub struct ApplyArgumentTangent<'a, R>
743    where
744        R: Runtime,
745    {
746        /// Parameter direction.
747        pub parameters: &'a ParameterDirection,
748        /// Input-state direction.
749        pub state: &'a R::Buffer,
750    }
751
752    impl<'a, R> Differentiable for ApplyArguments<'a, R>
753    where
754        R: Runtime,
755    {
756        type Tangent = ApplyArgumentTangent<'a, R>;
757        type Cotangent = ApplyCotangents<R::Buffer>;
758    }
759
760    /// Explicit rule object connecting QMBED native application to ChainRules.
761    pub struct ApplyRule<'a, R>
762    where
763        R: Runtime,
764    {
765        /// Execution runtime.
766        pub runtime: &'a R,
767        /// Parameterized operator family.
768        pub operator: &'a QuantumOperator,
769    }
770
771    impl<'rule, 'args, R> JvpRule<ApplyArguments<'args, R>> for ApplyRule<'rule, R>
772    where
773        R: Runtime,
774        Operator: RuntimeLinearOperator<R>,
775    {
776        type Output = State<R::Buffer>;
777        type Error = QmbedError;
778
779        fn jvp(
780            &self,
781            args: &ApplyArguments<'args, R>,
782            tangent: &ApplyArgumentTangent<'args, R>,
783        ) -> Result<(Self::Output, <Self::Output as Differentiable>::Tangent)> {
784            let result = apply_jvp(
785                self.runtime,
786                self.operator,
787                args.parameters,
788                tangent.parameters,
789                args.state,
790                tangent.state,
791            )?;
792            Ok((State(result.value), result.tangent))
793        }
794    }
795
796    /// ChainRules wrapper around QMBED's native one-shot pullback.
797    pub struct RulePullback<'a, R>
798    where
799        R: Runtime,
800    {
801        inner: ApplyPullback<'a, R>,
802    }
803
804    impl<R> Pullback<R::Buffer, ApplyCotangents<R::Buffer>> for RulePullback<'_, R>
805    where
806        R: Runtime,
807        Operator: RuntimeLinearOperator<R> + RuntimeAdjointLinearOperator<R>,
808    {
809        type Error = QmbedError;
810
811        fn apply(self, cotangent: R::Buffer) -> Result<ApplyCotangents<R::Buffer>> {
812            self.inner.backward(&cotangent)
813        }
814    }
815
816    impl<'rule, 'args, R> VjpRule<ApplyArguments<'args, R>> for ApplyRule<'rule, R>
817    where
818        R: Runtime,
819        Operator: RuntimeLinearOperator<R> + RuntimeAdjointLinearOperator<R>,
820    {
821        type Output = State<R::Buffer>;
822        type Error = QmbedError;
823        type Pullback<'a>
824            = RulePullback<'a, R>
825        where
826            Self: 'a,
827            ApplyArguments<'args, R>: 'a;
828
829        fn vjp<'a>(
830            &'a self,
831            args: &'a ApplyArguments<'args, R>,
832        ) -> Result<(Self::Output, Self::Pullback<'a>)> {
833            let (value, pullback) =
834                apply_vjp(self.runtime, self.operator, args.parameters, args.state)?;
835            Ok((State(value), RulePullback { inner: pullback }))
836        }
837    }
838
839    /// ChainRules protocol object for an isolated ground-state energy.
840    pub struct GroundStateEnergyRule<'a> {
841        /// Real-parameter Hermitian operator family.
842        pub operator: &'a QuantumOperator,
843        /// Solver and convergence controls, including two or more eigenpairs.
844        pub options: EigshOptions,
845    }
846
847    /// One-shot scalar-energy pullback owning the computed gradient.
848    pub struct GroundStateEnergyPullback {
849        gradient: ParameterGradient,
850    }
851
852    impl Pullback<f64, ParameterGradient> for GroundStateEnergyPullback {
853        type Error = QmbedError;
854
855        fn apply(self, cotangent: f64) -> Result<ParameterGradient> {
856            if !cotangent.is_finite() {
857                return Err(QmbedError::InvalidOptions(
858                    "ground-state energy cotangent must be finite".into(),
859                ));
860            }
861            Ok(ParameterGradient::new(
862                Arc::clone(self.gradient.schema()),
863                self.gradient
864                    .values()
865                    .iter()
866                    .map(|value| cotangent * *value)
867                    .collect(),
868            ))
869        }
870    }
871
872    impl JvpRule<ParameterValues> for GroundStateEnergyRule<'_> {
873        type Output = f64;
874        type Error = QmbedError;
875
876        fn jvp(
877            &self,
878            parameters: &ParameterValues,
879            tangent: &ParameterDirection,
880        ) -> Result<(f64, f64)> {
881            if parameters.schema().as_ref() != tangent.schema().as_ref() {
882                return Err(QmbedError::InvalidOptions(
883                    "ground-state tangent uses a different parameter schema".into(),
884                ));
885            }
886            let result = ground_state_energy_gradient(
887                self.operator,
888                parameters,
889                self.options.clone(),
890                &mut EigshWorkspace::new(),
891            )?;
892            if result.diagnostics.status != GradientStatus::Reliable {
893                return Err(QmbedError::InvalidOptions(format!(
894                    "ground-state derivative is not reliable: {:?}",
895                    result.diagnostics
896                )));
897            }
898            let directional = result
899                .gradient
900                .values()
901                .iter()
902                .zip(tangent.values())
903                .map(|(gradient, direction)| (direction.conj() * gradient).re)
904                .sum();
905            Ok((result.energy, directional))
906        }
907    }
908
909    impl VjpRule<ParameterValues> for GroundStateEnergyRule<'_> {
910        type Output = f64;
911        type Error = QmbedError;
912        type Pullback<'a>
913            = GroundStateEnergyPullback
914        where
915            Self: 'a,
916            ParameterValues: 'a;
917
918        fn vjp<'a>(&'a self, parameters: &'a ParameterValues) -> Result<(f64, Self::Pullback<'a>)> {
919            let result = ground_state_energy_gradient(
920                self.operator,
921                parameters,
922                self.options.clone(),
923                &mut EigshWorkspace::new(),
924            )?;
925            if result.diagnostics.status != GradientStatus::Reliable {
926                return Err(QmbedError::InvalidOptions(format!(
927                    "ground-state derivative is not reliable: {:?}",
928                    result.diagnostics
929                )));
930            }
931            Ok((
932                result.energy,
933                GroundStateEnergyPullback {
934                    gradient: result.gradient,
935                },
936            ))
937        }
938    }
939}
qmbed/ad/mod.rs

qmbed/ad/
mod.rs
