Rustb/
ndarray_lapack.rs

1//!这个模块是用来求解大矩阵的部分本征值的模块, 用的lapackc的 cheevx 等函数求解.
2#[cfg(any(feature = "intel-mkl-system", feature = "intel-mkl-static"))]
3extern crate intel_mkl_src as _src;
4
5#[cfg(any(feature = "openblas-system", feature = "openblas-static"))]
6extern crate openblas_src as _src;
7
8#[cfg(any(feature = "netlib-system", feature = "netlib-static"))]
9extern crate netlib_src as _src;
10
11use lapack::{cheevx, zheev, zheevr, zheevr_2stage, zheevx};
12use ndarray::{Array1, Array2, ArrayBase, Data, Ix2};
13use ndarray_linalg::EigValsh;
14use ndarray_linalg::UPLO;
15use num_complex::Complex;
16use std::ffi::c_char;
17
18pub fn eigh_x<S>(
19    x: &ArrayBase<S, Ix2>,
20    range: (f64, f64),
21    epsilon: f64,
22    uplo: UPLO,
23) -> (Array1<f64>, Array2<Complex<f64>>)
24where
25    S: Data<Elem = Complex<f64>>,
26{
27    // 获取矩阵的阶数
28    let n = x.shape()[0] as i32;
29    // 创建一个可变的副本，用于存储特征向量或被销毁
30    let mut a: Vec<_> = x.iter().cloned().collect();
31    // 创建一个可变的向量，用于存储特征值
32    let mut w = vec![0.0; n as usize];
33    // 创建一个可变的向量，用于存储特征向量
34    let mut z = vec![Complex::new(0.0, 0.0); (n * n) as usize];
35    // 创建一个可变的变量，用于存储特征值的个数
36    let mut m = 0;
37    // 创建一个可变的变量，用于存储函数的返回状态
38    let mut info = 0;
39    // 创建一个可变的向量，用于存储失败的特征值的索引
40    let mut ifail = vec![0; n as usize];
41    // 创建一个可变的向量，用于作为工作空间
42    let mut work = vec![Complex::new(0.0, 0.0); (2 * n) as usize];
43    // 创建一个可变的向量，用于作为实数工作空间
44    let mut rwork = vec![0.0; (7 * n) as usize];
45    // 创建一个可变的向量，用于作为整数工作空间
46    let mut iwork = vec![0; (5 * n) as usize];
47    // 调用 cheevx 函数，使用 'N' 表示不计算特征向量，使用 'V' 表示计算范围内的特征值，使用 'U' 表示输入的矩阵是上三角
48    // Assuming range is a tuple like (f64, f64)
49    let job1 = b'V';
50    let job2 = b'V';
51    let job3 = match uplo {
52        UPLO::Upper => b'U',
53        UPLO::Lower => b'L',
54    };
55
56    unsafe {
57        zheevx(
58            job1,
59            job2,
60            job3,
61            n,
62            &mut a,
63            n,
64            range.0,
65            range.1,
66            0,
67            n,
68            epsilon,
69            &mut m,
70            &mut w,
71            &mut z,
72            n,
73            &mut work,
74            2 * n,
75            &mut rwork,
76            &mut iwork,
77            &mut ifail,
78            &mut info,
79        );
80    }
81    // 检查函数的返回状态，如果是 0，表示成功，否则表示有问题
82    if info == 0 {
83        // 将特征值向量转换为一维数组并返回
84        (
85            Array1::<f64>::from_vec(w.into_iter().take(m as usize).collect()),
86            Array2::<Complex<f64>>::from_shape_vec(
87                [m as usize, n as usize],
88                z.into_iter().take((n * m) as usize).collect(),
89            )
90            .unwrap(),
91        )
92    } else {
93        // 报告错误信息
94        panic!("cheevx failed with info = {}", info);
95    }
96}
97
98pub fn eigvalsh_x<S>(
99    x: &ArrayBase<S, Ix2>,
100    range: (f64, f64),
101    epsilon: f64,
102    uplo: UPLO,
103) -> Array1<f64>
104where
105    S: Data<Elem = Complex<f64>>,
106{
107    // 获取矩阵的阶数
108    let n = x.shape()[0] as i32;
109    // 创建一个可变的副本，用于存储特征向量或被销毁
110    let mut a: Vec<_> = x.iter().cloned().collect();
111    // 创建一个可变的向量，用于存储特征值
112    let mut w = vec![0.0; n as usize];
113    // 创建一个可变的向量，用于存储特征向量
114    let mut z = vec![Complex::new(0.0, 0.0); (n * n) as usize];
115    // 创建一个可变的变量，用于存储特征值的个数
116    let mut m = 0;
117    // 创建一个可变的变量，用于存储函数的返回状态
118    let mut info = 0;
119    // 创建一个可变的向量，用于存储失败的特征值的索引
120    let mut ifail = vec![0; n as usize];
121    // 创建一个可变的向量，用于存储失败的特征值的索引
122    let mut work = vec![Complex::new(0.0, 0.0); (2 * n) as usize];
123    // 创建一个可变的向量，用于作为实数工作空间
124    let mut rwork = vec![0.0; (7 * n) as usize];
125    // 创建一个可变的向量，用于作为整数工作空间
126    let mut iwork = vec![0; (5 * n) as usize];
127    // 调用 cheevx 函数，使用 'N' 表示不计算特征向量，使用 'V' 表示计算范围内的特征值，使用 'U' 表示输入的矩阵是上三角
128    // Assuming range is a tuple like (f64, f64)
129    let job1 = b'N';
130    let job2 = b'V';
131    let job3 = match uplo {
132        UPLO::Upper => b'U',
133        UPLO::Lower => b'L',
134    };
135
136    unsafe {
137        zheevx(
138            job1,
139            job2,
140            job3,
141            n,
142            &mut a,
143            n,
144            range.0,
145            range.1,
146            0,
147            n,
148            epsilon,
149            &mut m,
150            &mut w,
151            &mut z,
152            n,
153            &mut work,
154            2 * n,
155            &mut rwork,
156            &mut iwork,
157            &mut ifail,
158            &mut info,
159        );
160    }
161    // 检查函数的返回状态，如果是 0，表示成功，否则表示有问题
162    if info == 0 {
163        // 将特征值向量转换为一维数组并返回
164        Array1::<f64>::from_vec(w.into_iter().take(m as usize).collect())
165    } else {
166        // 报告错误信息
167        panic!("cheevx failed with info = {}", info);
168    }
169}
170
171pub fn eigh_r<S>(
172    x: &ArrayBase<S, Ix2>,
173    range: (f64, f64),
174    epsilon: f64,
175    uplo: UPLO,
176) -> (Array1<f64>, Array2<Complex<f64>>)
177where
178    S: Data<Elem = Complex<f64>>,
179{
180    let job1 = b'V';
181    let job2 = b'V';
182    let job3 = match uplo {
183        UPLO::Upper => b'U',
184        UPLO::Lower => b'L',
185    };
186    // 获取矩阵的阶数
187    let n = x.shape()[0] as i32;
188    // 创建一个可变的副本，用于存储特征向量或被销毁
189    let mut a: Vec<_> = x.iter().cloned().collect();
190    // 创建一个可变的向量，用于存储特征值
191    let mut w = vec![0.0; n as usize];
192    // 创建一个可变的向量，用于存储特征向量
193    let mut z = vec![Complex::new(0.0, 0.0); (n * n) as usize];
194    let mut isuppz = vec![0; 2 * n as usize];
195    // 创建一个可变的变量，用于存储特征值的个数
196    let mut m = 0;
197    // 创建一个可变的变量，用于存储函数的返回状态
198    let mut info = 0;
199    /*
200    // 创建一个可变的向量，用于作为工作空间
201    let mut work = vec![Complex::new(0.0,0.0); 1 as usize];
202    // 创建一个可变的向量，用于作为实数工作空间
203    let mut rwork = vec![0.0; 1 as usize];
204    // 创建一个可变的向量，用于作为整数工作空间
205    let mut iwork = vec![0; 1 as usize];
206    // 调用 cheevx 函数，使用 'N' 表示不计算特征向量，使用 'V' 表示计算范围内的特征值，使用 'U' 表示输入的矩阵是上三角
207    // Assuming range is a tuple like (f64, f64)
208
209    unsafe {
210         zheevr(
211            job1,
212            job2,
213            job3,
214            n ,
215            &mut a,
216            n ,
217            range.0,
218            range.1,
219            0 ,
220            n ,
221            epsilon ,
222            &mut m,
223            &mut w,
224            &mut z,
225            n ,
226            &mut isuppz,
227            &mut work,
228            -1 ,
229            &mut rwork,
230            -1,
231            &mut iwork,
232            -1,
233            &mut info,
234        );
235    }
236
237    let lwork=work[0].re as i32;
238    let liwork=iwork[0] as i32;
239    let lrwork=rwork[0] as i32;
240    */
241    let lwork = n * 33 as i32;
242    let liwork = n * 10 as i32;
243    let lrwork = n * 24 as i32;
244    let mut work = vec![Complex::new(0.0, 0.0); lwork as usize];
245    // 创建一个可变的向量，用于作为实数工作空间
246    let mut rwork = vec![0.0; lrwork as usize];
247    // 创建一个可变的向量，用于作为整数工作空间
248    let mut iwork = vec![0; liwork as usize];
249    // 调用 cheevx 函数，使用 'N' 表示不计算特征向量，使用 'V' 表示计算范围内的特征值，使用 'U' 表示输入的矩阵是上三角
250
251    unsafe {
252        zheevr(
253            job1,
254            job2,
255            job3,
256            n,
257            &mut a,
258            n,
259            range.0,
260            range.1,
261            0,
262            n,
263            epsilon,
264            &mut m,
265            &mut w,
266            &mut z,
267            n,
268            &mut isuppz,
269            &mut work,
270            lwork,
271            &mut rwork,
272            lrwork,
273            &mut iwork,
274            liwork,
275            &mut info,
276        );
277    }
278
279    // 检查函数的返回状态，如果是 0，表示成功，否则表示有问题
280    if info == 0 {
281        // 将特征值向量转换为一维数组并返回
282        (
283            Array1::<f64>::from_vec(w.into_iter().take(m as usize).collect()),
284            Array2::<Complex<f64>>::from_shape_vec(
285                [m as usize, n as usize],
286                z.into_iter().take((n * m) as usize).collect(),
287            )
288            .unwrap(),
289        )
290    } else {
291        // 报告错误信息
292        panic!("cheevx failed with info = {}", info);
293    }
294}
295
296pub fn eigvalsh_r<S>(
297    x: &ArrayBase<S, Ix2>,
298    range: (f64, f64),
299    epsilon: f64,
300    uplo: UPLO,
301) -> Array1<f64>
302where
303    S: Data<Elem = Complex<f64>>,
304{
305    // 获取矩阵的阶数
306    let n = x.shape()[0] as i32;
307    // 创建一个可变的副本，用于存储特征向量或被销毁
308    let mut a: Vec<_> = x.iter().cloned().collect();
309    // 创建一个可变的向量，用于存储特征值
310    let mut w = vec![0.0; n as usize];
311    // 创建一个可变的向量，用于存储特征向量
312    let mut z = vec![Complex::new(0.0, 0.0); (n * n) as usize];
313    let mut isuppz = vec![0; 2 * n as usize];
314    // 创建一个可变的变量，用于存储特征值的个数
315    let mut m = 0;
316    // 创建一个可变的变量，用于存储函数的返回状态
317    let mut info = 0;
318    // 创建一个可变的向量，用于作为工作空间
319    let mut work = vec![Complex::new(0.0, 0.0); 1 as usize];
320    // 创建一个可变的向量，用于作为实数工作空间
321    let mut rwork = vec![0.0; 1 as usize];
322    // 创建一个可变的向量，用于作为整数工作空间
323    let mut iwork = vec![0; 1 as usize];
324    // 调用 cheevx 函数，使用 'N' 表示不计算特征向量，使用 'V' 表示计算范围内的特征值，使用 'U' 表示输入的矩阵是上三角
325    // Assuming range is a tuple like (f64, f64)
326    let job1 = b'N';
327    let job2 = b'V';
328    let job3 = match uplo {
329        UPLO::Upper => b'U',
330        UPLO::Lower => b'L',
331    };
332
333    unsafe {
334        zheevr(
335            job1,
336            job2,
337            job3,
338            n,
339            &mut a,
340            n,
341            range.0,
342            range.1,
343            0,
344            n,
345            epsilon,
346            &mut m,
347            &mut w,
348            &mut z,
349            n,
350            &mut isuppz,
351            &mut work,
352            -1,
353            &mut rwork,
354            -1,
355            &mut iwork,
356            -1,
357            &mut info,
358        );
359    }
360
361    let lwork = work[0].re as i32;
362    let liwork = iwork[0] as i32;
363    let lrwork = rwork[0] as i32;
364    let mut work = vec![Complex::new(0.0, 0.0); lwork as usize];
365    // 创建一个可变的向量，用于作为实数工作空间
366    let mut rwork = vec![0.0; lrwork as usize];
367    // 创建一个可变的向量，用于作为整数工作空间
368    let mut iwork = vec![0; liwork as usize];
369    // 调用 cheevx 函数，使用 'N' 表示不计算特征向量，使用 'V' 表示计算范围内的特征值，使用 'U' 表示输入的矩阵是上三角
370
371    unsafe {
372        zheevr(
373            job1,
374            job2,
375            job3,
376            n,
377            &mut a,
378            n,
379            range.0,
380            range.1,
381            0,
382            n,
383            epsilon,
384            &mut m,
385            &mut w,
386            &mut z,
387            n,
388            &mut isuppz,
389            &mut work,
390            lwork,
391            &mut rwork,
392            lrwork,
393            &mut iwork,
394            liwork,
395            &mut info,
396        );
397    }
398    // 检查函数的返回状态，如果是 0，表示成功，否则表示有问题
399    if info == 0 {
400        // 将特征值向量转换为一维数组并返回
401        Array1::<f64>::from_vec(w.into_iter().take(m as usize).collect())
402    } else {
403        // 报告错误信息
404        panic!("cheevx failed with info = {}", info);
405    }
406}
407
408pub fn eigvalsh_v<S>(x: &ArrayBase<S, Ix2>, uplo: UPLO) -> Array1<f64>
409where
410    S: Data<Elem = Complex<f64>>,
411{
412    // 获取矩阵的阶数
413    let n = x.shape()[0] as i32;
414    // 创建一个可变的副本，用于存储特征向量或被销毁
415    let mut a: Vec<_> = x.iter().cloned().collect();
416    // 创建一个可变的向量，用于存储特征值
417    let mut w = vec![0.0; n as usize];
418    // 创建一个可变的向量，用于存储特征向量
419    let mut z = vec![Complex::new(0.0, 0.0); (n * n) as usize];
420    let mut isuppz = vec![0; 2 * n as usize];
421    // 创建一个可变的变量，用于存储特征值的个数
422    let mut m = 0;
423    // 创建一个可变的变量，用于存储函数的返回状态
424    let mut info = 0;
425    // 创建一个可变的向量，用于作为工作空间
426    let mut work = vec![Complex::new(0.0, 0.0); 1 as usize];
427    // 创建一个可变的向量，用于作为实数工作空间
428    let mut rwork = vec![0.0; (3 * n - 2) as usize];
429    // 创建一个可变的向量，用于作为整数工作空间
430    let job1 = b'N';
431    let job2 = match uplo {
432        UPLO::Upper => b'U',
433        UPLO::Lower => b'L',
434    };
435
436    unsafe {
437        zheev(
438            job1, job2, n, &mut a, n, &mut w, &mut work, -1, &mut rwork, &mut info,
439        );
440    }
441    // 获取推荐的工作空间大小
442    let lwork = work[0].re as i32;
443    // 重新分配工作空间
444    work = vec![Complex::new(0.0, 0.0); lwork as usize];
445
446    // 第二次调用 zheev，实际计算特征值和特征向量
447    unsafe {
448        zheev(
449            job1, job2, n, &mut a, n, &mut w, &mut work, lwork, &mut rwork, &mut info,
450        );
451    }
452    // 检查函数的返回状态，如果是 0，表示成功，否则表示有问题
453    if info == 0 {
454        // 将特征值向量转换为一维数组并返回
455        Array1::<f64>::from_vec(w.into_iter().collect())
456    } else {
457        // 报告错误信息
458        panic!("zheev failed with info = {}", info);
459    }
460}
Rustb/ndarray_lapack.rs

Rustb/
ndarray_lapack.rs
