1#!/usr/bin/env python
2# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16# Author: Qiming Sun <osirpt.sun@gmail.com>
17#
18
19'''
20Generate symmetry adapted basis
21'''
22
23from functools import reduce
24import numpy
25from pyscf.data.elements import _symbol, _rm_digit
26from pyscf import gto
27from pyscf.lib.exceptions import PointGroupSymmetryError
28from pyscf.symm import geom
29from pyscf.symm import param
30
31__all__ = ['tot_parity_odd',
32           'symm_adapted_basis',
33           'dump_symm_adapted_basis',
34           'symmetrize_matrix',
35           'linearmole_symm_descent',
36           'linearmole_irrep_symb2id',
37           'linearmole_irrep_id2symb',
38           'linearmole_symm_adapted_basis',
39           'so3_irrep_symb2id',
40           'so3_irrep_id2symb',]
41
42OP_PARITY_ODD = {
43    'E'  : (0, 0, 0),
44    'C2x': (0, 1, 1),
45    'C2y': (1, 0, 1),
46    'C2z': (1, 1, 0),
47    'i'  : (1, 1, 1),
48    'sx' : (1, 0, 0),
49    'sy' : (0, 1, 0),
50    'sz' : (0, 0, 1),
51}
52
53def tot_parity_odd(op, l, m):
54    if op == 'E':
55        return 0
56    else:
57        ox,oy,oz = OP_PARITY_ODD[op]
58        gx,gy,gz = param.SPHERIC_GTO_PARITY_ODD[l][l+m]
59        return (ox and gx) ^ (oy and gy) ^ (oz and gz)
60
61def symm_adapted_basis(mol, gpname, orig=0, coordinates=None):
62    if gpname == 'SO3':
63        return so3_symm_adapted_basis(mol, gpname, orig, coordinates)
64    elif gpname in ('Dooh', 'Coov'):
65        return linearmole_symm_adapted_basis(mol, gpname, orig, coordinates)
66
67    # prop_atoms are the atoms relocated wrt the charge center with proper
68    # orientation
69    if coordinates is None:
70        coordinates = numpy.eye(3)
71    prop_atoms = mol.format_atom(mol._atom, orig, coordinates, 'Bohr')
72    eql_atom_ids = geom.symm_identical_atoms(gpname, prop_atoms)
73
74    ops = numpy.asarray([param.D2H_OPS[op] for op in param.OPERATOR_TABLE[gpname]])
75    chartab = numpy.array([x[1:] for x in param.CHARACTER_TABLE[gpname]])
76    nirrep = chartab.__len__()
77    aoslice = mol.aoslice_by_atom()
78    nao = mol.nao_nr()
79    atom_coords = numpy.array([a[1] for a in prop_atoms])
80
81    sodic = [[] for i in range(8)]
82    for atom_ids in eql_atom_ids:
83        r0 = atom_coords[atom_ids[0]]
84        op_coords = numpy.einsum('x,nxy->ny', r0, ops)
85# Using ops to generate other atoms from atom_ids[0]
86        coords0 = atom_coords[atom_ids]
87        dc = abs(op_coords.reshape(-1,1,3) - coords0).max(axis=2)
88        op_relate_idx = numpy.argwhere(dc < geom.TOLERANCE)[:,1]
89        ao_loc = numpy.array([aoslice[atom_ids[i],2] for i in op_relate_idx])
90
91        b0, b1 = aoslice[atom_ids[0],:2]
92        ip = 0
93        for ib in range(b0, b1):
94            l = mol.bas_angular(ib)
95            if mol.cart:
96                degen = (l + 1) * (l + 2) // 2
97                cbase = numpy.zeros((degen,nirrep,nao))
98                for op_id, op in enumerate(ops):
99                    n = 0
100                    for x in range(l, -1, -1):
101                        for y in range(l-x, -1, -1):
102                            z = l-x-y
103                            idx = ao_loc[op_id] + n
104                            sign = op[0,0]**x * op[1,1]**y * op[2,2]**z
105                            cbase[n,:,idx] += sign * chartab[:,op_id]
106                            n += 1
107            else:
108                degen = l * 2 + 1
109                cbase = numpy.zeros((degen,nirrep,nao))
110                for op_id, op in enumerate(param.OPERATOR_TABLE[gpname]):
111                    for n, m in enumerate(range(-l, l+1)):
112                        idx = ao_loc[op_id] + n
113                        if tot_parity_odd(op, l, m):
114                            cbase[n,:,idx] -= chartab[:,op_id]
115                        else:
116                            cbase[n,:,idx] += chartab[:,op_id]
117            norms = numpy.sqrt(numpy.einsum('mij,mij->mi', cbase, cbase))
118
119            for i in range(mol.bas_nctr(ib)):
120                for n, ir in numpy.argwhere(norms > 1e-12):
121                    c = numpy.zeros(nao)
122                    c[ip:] = cbase[n,ir,:nao-ip] / norms[n,ir]
123                    sodic[ir].append(c)
124                ip += degen
125
126    ao_loc = mol.ao_loc_nr()
127    l_idx = {}
128    ANG_OF = 1
129    for l in range(mol._bas[:,ANG_OF].max()+1):
130        idx = [numpy.arange(ao_loc[ib], ao_loc[ib+1])
131               for ib in numpy.where(mol._bas[:,ANG_OF] == l)[0]]
132        if idx:
133            l_idx[l] = numpy.hstack(idx)
134
135    Ds = _ao_rotation_matrices(mol, coordinates)
136    so = []
137    irrep_ids = []
138    for ir, c in enumerate(sodic):
139        if len(c) > 0:
140            irrep_ids.append(ir)
141            c_ir = numpy.vstack(c).T
142            nso = c_ir.shape[1]
143            for l, idx in l_idx.items():
144                c = c_ir[idx].reshape(-1,Ds[l].shape[1],nso)
145                c_ir[idx] = numpy.einsum('nm,smp->snp', Ds[l], c).reshape(-1,nso)
146
147            so.append(c_ir)
148
149    return so, irrep_ids
150
151def _ao_rotation_matrices(mol, axes):
152    '''Cache the rotation matrices'''
153    from pyscf import lib
154    from pyscf.symm.Dmatrix import Dmatrix, get_euler_angles
155    alpha, beta, gamma = get_euler_angles(numpy.eye(3), axes)
156    ANG_OF = 1
157    l_max = mol._bas[:,ANG_OF].max()
158
159    if not mol.cart:
160        return [Dmatrix(l, alpha, beta, gamma, reorder_p=True)
161                for l in range(l_max+1)]
162
163    pp = Dmatrix(1, alpha, beta, gamma, reorder_p=True)
164
165    Ds = [numpy.ones((1,1))]
166    for l in range(1, l_max+1):
167        # All possible x,y,z combinations
168        cidx = numpy.sort(lib.cartesian_prod([(0, 1, 2)] * l), axis=1)
169
170        addr = 0
171        affine = numpy.ones((1,1))
172        for i in range(l):
173            nd = affine.shape[0] * 3
174            affine = numpy.einsum('ik,jl->ijkl', affine, pp).reshape(nd, nd)
175            addr = addr * 3 + cidx[:,i]
176
177        uniq_addr, rev_addr = numpy.unique(addr, return_inverse=True)
178        ncart = (l + 1) * (l + 2) // 2
179        assert ncart == uniq_addr.size
180        trans = numpy.zeros((ncart,ncart))
181        for i, k in enumerate(rev_addr):
182            trans[k] += affine[i,uniq_addr]
183        Ds.append(trans)
184    return Ds
185
186def dump_symm_adapted_basis(mol, so):
187    raise PointGroupSymmetryError('TODO')
188
189def symmetrize_matrix(mat, so):
190    return [reduce(numpy.dot, (c.conj().T,mat,c)) for c in so]
191
192def _basis_offset_for_atoms(atoms, basis_tab):
193    basoff = [0]
194    n = 0
195    for at in atoms:
196        symb = _symbol(at[0])
197        if symb in basis_tab:
198            bas0 = basis_tab[symb]
199        else:
200            bas0 = basis_tab[_rm_digit(symb)]
201        for b in bas0:
202            angl = b[0]
203            n += _num_contract(b) * (angl*2+1)
204        basoff.append(n)
205    return n, basoff
206
207def _num_contract(basis):
208    if isinstance(basis[1], int):
209        # This branch should never be reached if basis_tab is formated by
210        # function mole.format_basis
211        nctr = len(basis[2]) - 1
212    else:
213        nctr = len(basis[1]) - 1
214    return nctr
215
216###############################
217# SO3 (real spherical harmonics)
218# Irreps ID maps
219# SO3       ->  Dooh
220# s       0 ->  A1g   0
221
222# pz    105 ->  A1u   5
223# py    106 ->  E1uy  6
224# px    107 ->  E1ux  7
225
226# dz2   200 ->  A1g   0
227# dyz   203 ->  E1gy  3
228# dxz   202 ->  E1gx  2
229# dxy   211 ->  E2gy  11
230# dx2y2 210 ->  E2gx  10
231
232# f0    305 ->  A1u   5
233# f-1   306 ->  E1uy  6
234# f+1   307 ->  E1ux  7
235# f-2   314 ->  E2uy  14
236# f+2   315 ->  E2ux  15
237# f-3   316 ->  E3uy  16
238# f+3   317 ->  E3ux  17
239
240# g0    400 ->  A1g   0
241# g-1   403 ->  E1gy  3
242# g+1   402 ->  E1gx  2
243# g-2   411 ->  E2gy  11
244# g+2   410 ->  E2gx  10
245# g-3   413 ->  E3gy  13
246# g+3   412 ->  E3gx  12
247# g-4   421 ->  E4gy  21
248# g+4   420 ->  E4gx  20
249_SO3_SYMB2ID = {
250    's+0' :     0,
251    'p-1':    106,
252    'p+0':    105,
253    'p+1':    107,
254    'd-2':    211,
255    'd-1':    203,
256    'd+0':    200,
257    'd+1':    202,
258    'd+2':    210,
259    'f-3':    316,
260    'f-2':    314,
261    'f-1':    306,
262    'f+0':    305,
263    'f+1':    307,
264    'f+2':    315,
265    'f+3':    317,
266    'g-4':    421,
267    'g-3':    413,
268    'g-2':    411,
269    'g-1':    403,
270    'g+0':    400,
271    'g+1':    402,
272    'g+2':    410,
273    'g+3':    412,
274    'g+4':    420,
275    'h-5':    526,
276    'h-4':    524,
277    'h-3':    516,
278    'h-2':    514,
279    'h-1':    506,
280    'h+0':    505,
281    'h+1':    507,
282    'h+2':    515,
283    'h+3':    517,
284    'h+4':    525,
285    'h+5':    527,
286    'i-6':    631,
287    'i-5':    623,
288    'i-4':    621,
289    'i-3':    613,
290    'i-2':    611,
291    'i-1':    603,
292    'i+0':    600,
293    'i+1':    602,
294    'i+2':    610,
295    'i+3':    612,
296    'i+4':    620,
297    'i+5':    622,
298    'i+6':    630,
299}
300_SO3_ID2SYMB = dict([(v, k) for k, v in _SO3_SYMB2ID.items()])
301_ANGULAR = 'spdfghik'
302
303def so3_irrep_symb2id(symb):
304    return _SO3_SYMB2ID[symb.lower()]
305
306def so3_irrep_id2symb(irrep_id):
307    return _SO3_ID2SYMB[irrep_id]
308
309def so3_symm_adapted_basis(mol, gpname, orig=0, coordinates=None):
310    assert gpname == 'SO3'
311    assert mol.natm == 1
312
313    ao_loc = mol.ao_loc_nr(cart=False)
314    nao_sph = ao_loc[-1]
315
316    if mol.cart:
317        coeff = mol.cart2sph(normalized='sp')
318    else:
319        coeff = numpy.eye(nao_sph)
320    nao = coeff.shape[0]
321
322    lmax = max(mol._bas[:,gto.ANG_OF])
323    so = []
324    irrep_names = []
325    for l in range(lmax+1):
326        bas_idx = mol._bas[:,gto.ANG_OF] == l
327        cs = [coeff[:,p0:p1]
328              for p0, p1 in zip(ao_loc[:-1][bas_idx], ao_loc[1:][bas_idx])]
329        c_groups = numpy.hstack(cs).reshape(nao, -1, l*2+1)
330        if l == 1:
331            so.extend([c_groups[:,:,1], c_groups[:,:,2], c_groups[:,:,0]])
332            irrep_names.extend(['p-1', 'p+0', 'p+1'])
333        else:
334            for m in range(-l, l+1):
335                so.append(c_groups[:,:,l+m])
336                irrep_names.append('%s%+d' % (_ANGULAR[l], m))
337
338    irrep_ids = [so3_irrep_symb2id(ir) for ir in irrep_names]
339    return so, irrep_ids
340
341
342###############################
343# Linear molecule
344# Irreps ID maps
345# Dooh     ->  D2h        |   Coov      -> C2v
346# A1g   0      Ag    0    |   A1    0      A1    0
347# A2g   1      B1g   1    |   A2    1      A2    1
348# A1u   5      B1u   5    |   E1x   2      B1    2
349# A2u   4      Au    4    |   E1y   3      B2    3
350# E1gx  2      B2g   2    |   E2x   10     A1    0
351# E1gy  3      B3g   3    |   E2y   11     A2    1
352# E1ux  7      B3u   7    |   E3x   12     B1    2
353# E1uy  6      B2u   6    |   E3y   13     B2    3
354# E2gx  10     Ag    0    |   E4x   20     A1    0
355# E2gy  11     B1g   1    |   E4y   21     A2    1
356# E2ux  15     B1u   5    |   E5x   22     B1    2
357# E2uy  14     Au    4    |   E5y   23     B2    3
358# E3gx  12     B2g   2    |
359# E3gy  13     B3g   3    |
360# E3ux  17     B3u   7    |
361# E3uy  16     B2u   6    |
362# E4gx  20     Ag    0    |
363# E4gy  21     B1g   1    |
364# E4ux  25     B1u   5    |
365# E4uy  24     Au    4    |
366# E5gx  22     B2g   2    |
367# E5gy  23     B3g   3    |
368# E5ux  27     B3u   7    |
369# E5uy  26     B2u   6    |
370
371DOOH_IRREP_ID_TABLE = {
372    'A1g' : 0,
373    'A2g' : 1,
374    'A1u' : 5,
375    'A2u' : 4,
376    'E1gx': 2,
377    'E1gy': 3,
378    'E1ux': 7,
379    'E1uy': 6,
380    '_evengx': 0,
381    '_evengy': 1,
382    '_evenux': 5,
383    '_evenuy': 4,
384    '_oddgx': 2,
385    '_oddgy': 3,
386    '_oddux': 7,
387    '_odduy': 6,
388}
389COOV_IRREP_ID_TABLE = {
390    'A1' : 0,
391    'A2' : 1,
392    'E1x': 2,
393    'E1y': 3,
394    '_evenx': 0,
395    '_eveny': 1,
396    '_oddx': 2,
397    '_oddy': 3,
398}
399
400def linearmole_symm_descent(gpname, irrepid):
401    '''Map irreps to D2h or C2v'''
402    if gpname in ('Dooh', 'Coov'):
403        return irrepid % 10
404    else:
405        raise PointGroupSymmetryError('%s is not proper for linear molecule.' % gpname)
406
407def linearmole_irrep_symb2id(gpname, symb):
408    if gpname == 'Dooh':
409        if symb in DOOH_IRREP_ID_TABLE:
410            return DOOH_IRREP_ID_TABLE[symb]
411        else:
412            n = int(''.join([i for i in symb if i.isdigit()]))
413            if n % 2:
414                return (n//2)*10 + DOOH_IRREP_ID_TABLE['_odd'+symb[-2:]]
415            else:
416                return (n//2)*10 + DOOH_IRREP_ID_TABLE['_even'+symb[-2:]]
417    elif gpname == 'Coov':
418        if symb in COOV_IRREP_ID_TABLE:
419            return COOV_IRREP_ID_TABLE[symb]
420        else:
421            n = int(''.join([i for i in symb if i.isdigit()]))
422            if n % 2:
423                return (n//2)*10 + COOV_IRREP_ID_TABLE['_odd'+symb[-1]]
424            else:
425                return (n//2)*10 + COOV_IRREP_ID_TABLE['_even'+symb[-1]]
426    else:
427        raise PointGroupSymmetryError('%s is not proper for linear molecule.' % gpname)
428
429DOOH_IRREP_SYMBS = ('A1g' , 'A2g' , 'E1gx', 'E1gy' , 'A2u', 'A1u' , 'E1uy', 'E1ux')
430DOOH_IRREP_SYMBS_EXT = ('gx' , 'gy' , 'gx', 'gy' , 'uy', 'ux' , 'uy', 'ux')
431COOV_IRREP_SYMBS = ('A1' , 'A2' , 'E1x', 'E1y')
432def linearmole_irrep_id2symb(gpname, irrep_id):
433    if gpname == 'Dooh':
434        if irrep_id < 10:
435            return DOOH_IRREP_SYMBS[irrep_id]
436        else:
437            n = irrep_id % 10
438            m = irrep_id // 10
439            if n in (0, 1, 5, 4):
440                rn = m*2
441            else:
442                rn = m*2+1
443            return 'E%d%s' % (rn, DOOH_IRREP_SYMBS_EXT[n])
444    elif gpname == 'Coov':
445        if irrep_id < 10:
446            return COOV_IRREP_SYMBS[irrep_id]
447        else:
448            n = irrep_id % 10
449            m = irrep_id // 10
450            if n < 2:
451                rn = m*2
452            else:
453                rn = m*2+1
454            if n % 2:
455                xy = 'y'
456            else:
457                xy = 'x'
458            return 'E%d%s' % (rn, xy)
459    else:
460        raise PointGroupSymmetryError('%s is not proper for linear molecule.' % gpname)
461
462def linearmole_symm_adapted_basis(mol, gpname, orig=0, coordinates=None):
463    assert(gpname in ('Dooh', 'Coov'))
464    assert(not mol.cart)
465
466    if coordinates is None:
467        coordinates = numpy.eye(3)
468    prop_atoms = mol.format_atom(mol._atom, orig, coordinates, 'Bohr')
469    eql_atom_ids = geom.symm_identical_atoms(gpname, prop_atoms)
470
471    aoslice = mol.aoslice_by_atom()
472    basoff = aoslice[:,2]
473    nao = mol.nao_nr()
474    sodic = {}
475    shalf = numpy.sqrt(.5)
476    def plus(i0, i1):
477        c = numpy.zeros(nao)
478        c[i0] = c[i1] = shalf
479        return c
480    def minus(i0, i1):
481        c = numpy.zeros(nao)
482        c[i0] = shalf
483        c[i1] =-shalf
484        return c
485    def identity(i0):
486        c = numpy.zeros(nao)
487        c[i0] = 1
488        return c
489    def add_so(irrep_name, c):
490        if irrep_name in sodic:
491            sodic[irrep_name].append(c)
492        else:
493            sodic[irrep_name] = [c]
494
495    if gpname == 'Dooh':
496        for atom_ids in eql_atom_ids:
497            if len(atom_ids) == 2:
498                at0 = atom_ids[0]
499                at1 = atom_ids[1]
500                ip = 0
501                b0, b1, p0, p1 = aoslice[at0]
502                for ib in range(b0, b1):
503                    angl = mol.bas_angular(ib)
504                    nc = mol.bas_nctr(ib)
505                    degen = angl * 2 + 1
506                    if angl == 1:
507                        for i in range(nc):
508                            aoff = ip + i*degen + angl
509# m = 0
510                            idx0 = basoff[at0] + aoff + 1
511                            idx1 = basoff[at1] + aoff + 1
512                            add_so('A1g', minus(idx0, idx1))
513                            add_so('A1u', plus (idx0, idx1))
514# m = +/- 1
515                            idx0 = basoff[at0] + aoff - 1
516                            idy0 = basoff[at0] + aoff
517                            idx1 = basoff[at1] + aoff - 1
518                            idy1 = basoff[at1] + aoff
519                            add_so('E1ux', plus (idx0, idx1))
520                            add_so('E1uy', plus (idy0, idy1))
521                            add_so('E1gx', minus(idx0, idx1))
522                            add_so('E1gy', minus(idy0, idy1))
523                    else:
524                        for i in range(nc):
525                            aoff = ip + i*degen + angl
526# m = 0
527                            idx0 = basoff[at0] + aoff
528                            idx1 = basoff[at1] + aoff
529                            if angl % 2: # p-sigma, f-sigma
530                                add_so('A1g', minus(idx0, idx1))
531                                add_so('A1u', plus (idx0, idx1))
532                            else: # s-sigma, d-sigma
533                                add_so('A1g', plus (idx0, idx1))
534                                add_so('A1u', minus(idx0, idx1))
535# +/-m
536                            for m in range(1,angl+1):
537                                idx0 = basoff[at0] + aoff + m
538                                idy0 = basoff[at0] + aoff - m
539                                idx1 = basoff[at1] + aoff + m
540                                idy1 = basoff[at1] + aoff - m
541                                if angl % 2: # odd parity
542                                    add_so('E%dux'%m, plus (idx0, idx1))
543                                    add_so('E%duy'%m, plus (idy0, idy1))
544                                    add_so('E%dgx'%m, minus(idx0, idx1))
545                                    add_so('E%dgy'%m, minus(idy0, idy1))
546                                else:
547                                    add_so('E%dgy'%m, plus (idy0, idy1))
548                                    add_so('E%dgx'%m, plus (idx0, idx1))
549                                    add_so('E%duy'%m, minus(idy0, idy1))
550                                    add_so('E%dux'%m, minus(idx0, idx1))
551                    ip += nc * degen
552            elif len(atom_ids) == 1:
553                at0 = atom_ids[0]
554                ip = 0
555                b0, b1, p0, p1 = aoslice[at0]
556                for ib in range(b0, b1):
557                    angl = mol.bas_angular(ib)
558                    nc = mol.bas_nctr(ib)
559                    degen = angl * 2 + 1
560                    if angl == 1:
561                        for i in range(nc):
562                            aoff = ip + i*degen + angl
563# m = 0
564                            idx0 = basoff[at0] + aoff + 1
565                            add_so('A1u', identity(idx0))
566# m = +/- 1
567                            idx0 = basoff[at0] + aoff - 1
568                            idy0 = basoff[at0] + aoff
569                            add_so('E1uy', identity(idy0))
570                            add_so('E1ux', identity(idx0))
571                    else:
572                        for i in range(nc):
573                            aoff = ip + i*degen + angl
574                            idx0 = basoff[at0] + aoff
575# m = 0
576                            if angl % 2:
577                                add_so('A1u', identity(idx0))
578                            else:
579                                add_so('A1g', identity(idx0))
580# +/-m
581                            for m in range(1,angl+1):
582                                idx0 = basoff[at0] + aoff + m
583                                idy0 = basoff[at0] + aoff - m
584                                if angl % 2: # p, f functions
585                                    add_so('E%dux'%m, identity(idx0))
586                                    add_so('E%duy'%m, identity(idy0))
587                                else: # d, g functions
588                                    add_so('E%dgy'%m, identity(idy0))
589                                    add_so('E%dgx'%m, identity(idx0))
590                    ip += nc * degen
591    elif gpname == 'Coov':
592        for atom_ids in eql_atom_ids:
593            at0 = atom_ids[0]
594            ip = 0
595            b0, b1, p0, p1 = aoslice[at0]
596            for ib in range(b0, b1):
597                angl = mol.bas_angular(ib)
598                nc = mol.bas_nctr(ib)
599                degen = angl * 2 + 1
600                if angl == 1:
601                    for i in range(nc):
602                        aoff = ip + i*degen + angl
603# m = 0
604                        idx0 = basoff[at0] + aoff + 1
605                        add_so('A1', identity(idx0))
606# m = +/- 1
607                        idx0 = basoff[at0] + aoff - 1
608                        idy0 = basoff[at0] + aoff
609                        add_so('E1x', identity(idx0))
610                        add_so('E1y', identity(idy0))
611                else:
612                    for i in range(nc):
613                        aoff = ip + i*degen + angl
614                        idx0 = basoff[at0] + aoff
615# m = 0
616                        add_so('A1', identity(idx0))
617# +/-m
618                        for m in range(1,angl+1):
619                            idx0 = basoff[at0] + aoff + m
620                            idy0 = basoff[at0] + aoff - m
621                            add_so('E%dx'%m, identity(idx0))
622                            add_so('E%dy'%m, identity(idy0))
623                ip += nc * degen
624
625    irrep_ids = []
626    irrep_names = list(sodic.keys())
627    for irname in irrep_names:
628        irrep_ids.append(linearmole_irrep_symb2id(gpname, irname))
629    irrep_idx = numpy.argsort(irrep_ids)
630    irrep_ids = [irrep_ids[i] for i in irrep_idx]
631
632    ao_loc = mol.ao_loc_nr()
633    l_idx = {}
634    ANG_OF = 1
635    for l in range(mol._bas[:,ANG_OF].max()+1):
636        idx = [numpy.arange(ao_loc[ib], ao_loc[ib+1])
637               for ib in numpy.where(mol._bas[:,ANG_OF] == l)[0]]
638        if idx:
639            l_idx[l] = numpy.hstack(idx)
640
641    Ds = _ao_rotation_matrices(mol, coordinates)
642    so = []
643    for i in irrep_idx:
644        c_ir = numpy.vstack(sodic[irrep_names[i]]).T
645        nso = c_ir.shape[1]
646        for l, idx in l_idx.items():
647            c = c_ir[idx].reshape(-1,Ds[l].shape[1],nso)
648            c_ir[idx] = numpy.einsum('nm,smp->snp', Ds[l], c).reshape(-1,nso)
649
650        so.append(c_ir)
651
652    return so, irrep_ids
653
654
655if __name__ == "__main__":
656    h2o = gto.Mole()
657    h2o.verbose = 0
658    h2o.output = None
659    h2o.atom = [['O' , (1. , 0.    , 0.   ,)],
660                [1   , (0. , -.757 , 0.587,)],
661                [1   , (0. , 0.757 , 0.587,)] ]
662    h2o.basis = {'H': 'cc-pvdz',
663                 'O': 'cc-pvdz',}
664    h2o.build()
665    gpname, origin, axes = geom.detect_symm(h2o._atom)
666    atoms = gto.format_atom(h2o._atom, origin, axes)
667    h2o.build(False, False, atom=atoms)
668    print(gpname)
669    eql_atoms = geom.symm_identical_atoms(gpname, atoms)
670    print(symm_adapted_basis(h2o, gpname, eql_atoms)[1])
671
672    mol = gto.M(
673        atom = [['H', (0,0,0)], ['H', (0,0,-1)], ['H', (0,0,1)]],
674        basis = 'ccpvtz', charge=1)
675    gpname, orig, axes = geom.detect_symm(mol._atom)
676    atoms = gto.format_atom(mol._atom, orig, axes)
677    mol.build(False, False, atom=atoms)
678    print(gpname)
679    eql_atoms = geom.symm_identical_atoms(gpname, atoms)
680    print(symm_adapted_basis(mol, gpname, eql_atoms)[1])
681
682    mol = gto.M(
683        atom = [['H', (0,0,0)], ['H', (0,0,-1)], ['He', (0,0,1)]],
684        basis = 'ccpvtz')
685    gpname, orig, axes = geom.detect_symm(mol._atom)
686    atoms = gto.format_atom(mol._atom, orig, axes)
687    mol.build(False, False, atom=atoms)
688    print(gpname)
689    eql_atoms = geom.symm_identical_atoms(gpname, atoms)
690    print(symm_adapted_basis(mol, gpname, eql_atoms)[1])
691