qudit_inst/numerical/functions/
hs.rs

1use std::sync::Arc;
2
3use super::super::InstantiationProblem;
4use crate::InstantiationTarget;
5use crate::numerical::Hessian;
6use crate::numerical::Jacobian;
7use crate::numerical::ProvidesJacobian;
8use crate::numerical::ProvidesResidualFunction;
9use crate::numerical::ResidualFunction;
10use crate::numerical::function::CostFunction;
11use crate::numerical::function::Function;
12use crate::numerical::function::Gradient;
13use crate::numerical::problem::Problem;
14use crate::numerical::problem::ProvidesCostFunction;
15use faer::ColMut;
16use faer::MatMut;
17use faer::RowMut;
18use faer::reborrow::ReborrowMut;
19use qudit_circuit::QuditCircuit;
20use qudit_core::ComplexScalar;
21use qudit_core::RealScalar;
22use qudit_expr::DifferentiationLevel;
23use qudit_expr::FUNCTION;
24use qudit_expr::GRADIENT;
25// use qudit_expr::HESSIAN;
26use qudit_tensor::PinnedTNVM;
27use qudit_tensor::TNVM;
28use qudit_tensor::compile_network;
29
30#[derive(Clone)]
31pub struct HSProblem<R: RealScalar> {
32    pub circuit: Arc<QuditCircuit>,
33    pub target: Arc<InstantiationTarget<R::C>>,
34}
35
36impl<R: RealScalar> HSProblem<R> {
37    pub fn new(circuit: Arc<QuditCircuit>, target: Arc<InstantiationTarget<R::C>>) -> Self {
38        HSProblem { circuit, target }
39    }
40
41    pub fn build_cost<const D: DifferentiationLevel>(&self) -> PinnedTNVM<R::C, D> {
42        let mut builder = self.circuit.as_tensor_network_builder();
43        // TODO: assert target.num_outputs = builder.num_open_outputs
44        // TODO: assert target.num_inputs = builder.num_open_inputs
45        // TODO: assert target.num_batch = builder.num_batch || target.num_batch = 1
46        match &*self.target {
47            InstantiationTarget::UnitaryMatrix(u) => {
48                let qudits = builder.open_output_indices();
49                builder = builder.prepend_unitary(u.dagger(), qudits);
50            }
51        }
52
53        let code = compile_network(builder.trace_all_open_wires().build());
54        TNVM::<R::C, D>::new(&code, Some(&self.circuit.params().const_map()))
55    }
56
57    pub fn build_residual<const D: DifferentiationLevel>(&self) -> PinnedTNVM<R::C, D> {
58        // TODO: deduplicate code
59        let mut builder = self.circuit.as_tensor_network_builder();
60        // TODO: assert target.num_outputs = builder.num_open_outputs
61        // TODO: assert target.num_inputs = builder.num_open_inputs
62        // TODO: assert target.num_batch = builder.num_batch || target.num_batch = 1
63        match &*self.target {
64            InstantiationTarget::UnitaryMatrix(u) => {
65                let qudits = builder.open_output_indices();
66                builder = builder.prepend_unitary(u.dagger(), qudits);
67            }
68        }
69
70        let code = compile_network(builder.build());
71        // dbg!(&code);
72        TNVM::<R::C, D>::new(&code, Some(&self.circuit.params().const_map()))
73    }
74}
75
76impl<R: RealScalar> InstantiationProblem<R> for HSProblem<R> {
77    fn from_instantiation(
78        circuit: Arc<QuditCircuit>,
79        target: Arc<InstantiationTarget<R::C>>,
80        _data: Arc<crate::DataMap>,
81    ) -> Self {
82        Self::new(circuit, target)
83    }
84}
85
86impl<R: RealScalar> Problem for HSProblem<R> {
87    fn num_params(&self) -> usize {
88        self.circuit.num_params()
89    }
90}
91
92impl<R: RealScalar> ProvidesCostFunction<R> for HSProblem<R> {
93    type CostFunction = HSFunction<R, 0>;
94
95    fn build_cost_function(&self) -> Self::CostFunction {
96        HSFunction {
97            tnvm: self.build_cost(),
98            n: R::from64(4.0),
99        } // TODO: N comes from circuit
100    }
101}
102
103impl<R: RealScalar> ProvidesResidualFunction<R> for HSProblem<R> {
104    type ResidualFunction = HSFunction<R, 0>;
105
106    fn build_residual_function(&self) -> Self::ResidualFunction {
107        HSFunction {
108            tnvm: self.build_residual(),
109            n: R::from64(4.0),
110        }
111    }
112}
113
114impl<R: RealScalar> ProvidesJacobian<R> for HSProblem<R> {
115    type Jacobian = HSFunction<R, GRADIENT>;
116
117    fn build_jacobian(&self) -> Self::Jacobian {
118        HSFunction {
119            tnvm: self.build_residual(),
120            n: R::from64(4.0),
121        }
122    }
123}
124
125pub struct HSFunction<R: RealScalar, const D: DifferentiationLevel> {
126    tnvm: PinnedTNVM<R::C, D>,
127    n: R,
128}
129
130impl<R: RealScalar, const D: DifferentiationLevel> Function for HSFunction<R, D> {
131    fn num_params(&self) -> usize {
132        self.tnvm.num_params()
133        // self.tnvm.borrow().num_params()
134    }
135}
136
137use num_complex::ComplexFloat;
138
139impl<R: RealScalar, const D: DifferentiationLevel> CostFunction<R> for HSFunction<R, D> {
140    fn cost(&mut self, params: &[R]) -> R {
141        let result = self.tnvm.evaluate::<FUNCTION>(params);
142        let trace = result.get_fn_result().unpack_scalar(); // This isn't a scalar if kraus
143        // dimension // TODO: add unpack vector able to handle scalar
144        let inner = trace.abs() / self.n;
145        R::from64(1.0) - inner
146        // (1 - (trace.clone().abs() / self.N).square()).sqrt()
147        // TODO: consider using proper norm (might be faster to use for gradient and hessian due
148        // not having to compute abs inside for loop, i.e. trade one sqrt for num_params.
149    }
150}
151
152// TODO: some safety so D>=1
153impl<R: RealScalar, const D: DifferentiationLevel> Gradient<R> for HSFunction<R, D> {
154    fn gradient_into(&mut self, params: &[R], grad_out: RowMut<R>) {
155        self.cost_and_gradient_into(params, grad_out);
156    }
157
158    fn cost_and_gradient_into(&mut self, params: &[R], grad_out: RowMut<R>) -> R {
159        let result = self.tnvm.evaluate::<GRADIENT>(params);
160        let grad_trace = result.get_grad_result().unpack_vector();
161        let trace = result.get_fn_result().unpack_scalar();
162        let trace_re = trace.real();
163        let trace_im = trace.imag();
164        for (out, trace) in grad_out.iter_mut().zip(grad_trace.iter()) {
165            let trace_re_d = trace.real();
166            let trace_im_d = trace.imag();
167            let num = trace_re * trace_re_d + trace_im * trace_im_d;
168            let dem = self.n * trace.abs();
169            *out = -(num / dem)
170        }
171        R::from64(1.0) - (trace.abs() / self.n)
172    }
173}
174
175impl<R: RealScalar, const D: DifferentiationLevel> Hessian<R> for HSFunction<R, D> {
176    fn hessian_into(&mut self, _params: &[R], _hess_out: faer::MatMut<R>) {
177        todo!()
178    }
179
180    fn cost_gradient_and_hessian_into(
181        &mut self,
182        _params: &[R],
183        _grad_out: RowMut<R>,
184        _hess_out: faer::MatMut<R>,
185    ) -> R {
186        // let result = self.tnvm.evaluate::<HESSIAN>(params);
187        // let hess_trace = result.get_hess_result().unpack_symsq_matrix();
188        // let grad_trace = result.get_grad_result().unpack_vector();
189        // let trace = result.get_fn_result().unpack_scalar();
190
191        // for (hess_out_col, (hess_trace_x_col, grad_trace_x)) in hess_out.col_iter_mut().zip(hess_trace.col_iter()
192
193        todo!("Need better iterators")
194        // Math = (-1/|trace|^3)*[R^3R_xy + RR_xyI^2 + I^2R_yR_x - RR_yII_x - II_yRR_x + R^2I_yI_x
195        // + R^2II_xy + I^3I_xy]
196        // R = Re(trace)
197        // R_x = Re(grad_trace @ x)
198        // Re_xy = Re(hess_trace @ xy)
199        // I = Im..
200
201        // R::new(1.0) - (trace.abs() / self.N)
202    }
203}
204
205impl<R: RealScalar, const D: DifferentiationLevel> ResidualFunction<R> for HSFunction<R, D> {
206    fn num_residuals(&self) -> usize {
207        let out_shape = self.tnvm.out_shape();
208        (out_shape.num_elements() - out_shape.nmats()) * 2
209    }
210
211    fn residuals_into(&mut self, params: &[R], mut residuals_out: ColMut<R>) {
212        let result = self.tnvm.evaluate::<FUNCTION>(params);
213        let kraus_ops = result.get_fn_result2().unpack_tensor3d();
214
215        let mut residual_index = 0;
216        for k in 0..kraus_ops.dims()[0] {
217            let mat = kraus_ops.subtensor_ref(k);
218
219            for (j, col) in mat.col_iter().enumerate() {
220                for (i, elem) in col.iter().enumerate() {
221                    if i == j {
222                        continue;
223                    }
224                    unsafe {
225                        let out = residuals_out.rb_mut().get_mut_unchecked(residual_index);
226                        *out = elem.real();
227                        residual_index += 1;
228                        let out = residuals_out.rb_mut().get_mut_unchecked(residual_index);
229                        *out = elem.imag();
230                        residual_index += 1;
231                    }
232                }
233            }
234
235            for diag_iter in 0..(mat.nrows() - 1) {
236                unsafe {
237                    let d_i = mat.get_unchecked(diag_iter, diag_iter);
238                    let d_j = mat.get_unchecked(diag_iter + 1, diag_iter + 1);
239                    let out = residuals_out.rb_mut().get_mut_unchecked(residual_index);
240                    *out = d_j.real() - d_i.real();
241                    residual_index += 1;
242                    let out = residuals_out.rb_mut().get_mut_unchecked(residual_index);
243                    *out = d_j.imag() - d_i.imag();
244                    residual_index += 1;
245                }
246            }
247        }
248    }
249}
250
251impl<R: RealScalar, const D: DifferentiationLevel> Jacobian<R> for HSFunction<R, D> {
252    fn residuals_and_jacobian_into(
253        &mut self,
254        params: &[R],
255        mut residuals_out: ColMut<R>,
256        mut jacobian_out: MatMut<R>,
257    ) {
258        let result = self.tnvm.evaluate::<GRADIENT>(params);
259        let kraus_ops = result.get_fn_result2().unpack_tensor3d();
260
261        // for k in 0..kraus_ops.dims()[0] {
262        //     let mat = kraus_ops.subtensor_ref(k);
263        //     for (j, col) in mat.col_iter().enumerate() {
264        //         for (i, elem) in col.iter().enumerate() {
265        //             print!("{} ", elem);
266        //         }
267        //         println!("");
268        //     }
269        //     println!("");
270        // }
271        // println!("");
272
273        let mut residual_index = 0;
274        for k in 0..kraus_ops.dims()[0] {
275            let mat = kraus_ops.subtensor_ref(k);
276
277            for (j, col) in mat.col_iter().enumerate() {
278                for (i, elem) in col.iter().enumerate() {
279                    if i == j {
280                        continue;
281                    }
282                    unsafe {
283                        let out = residuals_out.rb_mut().get_mut_unchecked(residual_index);
284                        *out = elem.real();
285                        residual_index += 1;
286                        let out = residuals_out.rb_mut().get_mut_unchecked(residual_index);
287                        *out = elem.imag();
288                        residual_index += 1;
289                    }
290                }
291            }
292
293            for diag_iter in 0..(mat.nrows() - 1) {
294                unsafe {
295                    let d_i = mat.get_unchecked(diag_iter, diag_iter);
296                    let d_j = mat.get_unchecked(diag_iter + 1, diag_iter + 1);
297                    let out = residuals_out.rb_mut().get_mut_unchecked(residual_index);
298                    *out = d_j.real() - d_i.real();
299                    residual_index += 1;
300                    let out = residuals_out.rb_mut().get_mut_unchecked(residual_index);
301                    *out = d_j.imag() - d_i.imag();
302                    residual_index += 1;
303                }
304            }
305        }
306
307        let grad_out = result.get_grad_result2().unpack_tensor4d();
308
309        for p in 0..grad_out.dims()[0] {
310            let partial_kraus_ops = grad_out.subtensor_ref(p);
311            let mut jacobian_col_out = jacobian_out.rb_mut().col_mut(p);
312            let mut residual_index = 0;
313
314            for k in 0..partial_kraus_ops.dims()[0] {
315                let mat = partial_kraus_ops.subtensor_ref(k);
316                // dbg!(&mat);
317
318                for (j, col) in mat.col_iter().enumerate() {
319                    for (i, elem) in col.iter().enumerate() {
320                        if i == j {
321                            continue;
322                        }
323                        unsafe {
324                            let out = jacobian_col_out.rb_mut().get_mut_unchecked(residual_index);
325                            *out = elem.real();
326                            residual_index += 1;
327                            let out = jacobian_col_out.rb_mut().get_mut_unchecked(residual_index);
328                            *out = elem.imag();
329                            residual_index += 1;
330                        }
331                    }
332                }
333
334                for diag_iter in 0..(mat.nrows() - 1) {
335                    unsafe {
336                        let d_i = mat.get_unchecked(diag_iter, diag_iter);
337                        let d_j = mat.get_unchecked(diag_iter + 1, diag_iter + 1);
338                        let out = jacobian_col_out.rb_mut().get_mut_unchecked(residual_index);
339                        *out = d_j.real() - d_i.real();
340                        residual_index += 1;
341                        let out = jacobian_col_out.rb_mut().get_mut_unchecked(residual_index);
342                        *out = d_j.imag() - d_i.imag();
343                        residual_index += 1;
344                    }
345                }
346            }
347        }
348    }
349
350    fn jacobian_into(&mut self, params: &[R], mut jacobian_out: MatMut<R>) {
351        let result = self.tnvm.evaluate::<GRADIENT>(params);
352        let grad_out = result.get_grad_result2().unpack_tensor4d();
353
354        for p in 0..grad_out.dims()[0] {
355            let partial_kraus_ops = grad_out.subtensor_ref(p);
356            let mut jacobian_col_out = jacobian_out.rb_mut().col_mut(p);
357            let mut residual_index = 0;
358
359            for k in 0..partial_kraus_ops.dims()[0] {
360                let mat = partial_kraus_ops.subtensor_ref(k);
361
362                for (j, col) in mat.col_iter().enumerate() {
363                    for (i, elem) in col.iter().enumerate() {
364                        if i == j {
365                            continue;
366                        }
367                        unsafe {
368                            let out = jacobian_col_out.rb_mut().get_mut_unchecked(residual_index);
369                            *out = elem.real();
370                            residual_index += 1;
371                            let out = jacobian_col_out.rb_mut().get_mut_unchecked(residual_index);
372                            *out = elem.imag();
373                            residual_index += 1;
374                        }
375                    }
376                }
377
378                for diag_iter in 0..(mat.nrows() - 1) {
379                    unsafe {
380                        let d_i = mat.get_unchecked(diag_iter, diag_iter);
381                        let d_j = mat.get_unchecked(diag_iter + 1, diag_iter + 1);
382                        let out = jacobian_col_out.rb_mut().get_mut_unchecked(residual_index);
383                        *out = d_j.real() - d_i.real();
384                        residual_index += 1;
385                        let out = jacobian_col_out.rb_mut().get_mut_unchecked(residual_index);
386                        *out = d_j.imag() - d_i.imag();
387                        residual_index += 1;
388                    }
389                }
390            }
391        }
392    }
393}
qudit_inst/numerical/functions/hs.rs

qudit_inst/numerical/functions/
hs.rs
