deploy: db8263b

_modules/ai2_kit/feat/spectrum/viber.html

@@ -606,8 +606,15 @@ <h1>Source code for ai2_kit.feat.spectrum.viber</h1><div class="highlight"><pre>
     <span class="c1"># get cell</span>
     <span class="n">cell</span> <span class="o">=</span> <span class="n">dp_sys</span><span class="o">.</span><span class="n">data</span><span class="p">[</span><span class="s1">&#39;cells&#39;</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>
 
+    <span class="c1"># get selected atoms ids</span>
+    <span class="n">symbols</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">dp_sys</span><span class="o">.</span><span class="n">data</span><span class="p">[</span><span class="s2">&quot;atom_names&quot;</span><span class="p">])[</span><span class="n">dp_sys</span><span class="o">.</span><span class="n">data</span><span class="p">[</span><span class="s2">&quot;atom_types&quot;</span><span class="p">]]</span>
+    <span class="n">sel_ids</span> <span class="o">=</span> <span class="p">[</span><span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">symbols</span> <span class="o">==</span> <span class="n">atype</span><span class="p">)[</span><span class="mi">0</span><span class="p">]</span> <span class="k">for</span> <span class="n">atype</span> <span class="ow">in</span> <span class="n">lumped_dict</span><span class="o">.</span><span class="n">keys</span><span class="p">()]</span>
+    <span class="n">sel_ids</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">concatenate</span><span class="p">(</span><span class="n">sel_ids</span><span class="p">)</span>
+
     <span class="c1"># build the data of atomic_dipole and atomic_polarizability with numpy</span>
     <span class="n">wannier_atoms</span> <span class="o">=</span> <span class="n">ase</span><span class="o">.</span><span class="n">io</span><span class="o">.</span><span class="n">read</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">data_dir</span><span class="p">,</span> <span class="n">wannier</span><span class="p">),</span> <span class="n">index</span><span class="o">=</span><span class="s2">&quot;:&quot;</span><span class="p">,</span> <span class="nb">format</span><span class="o">=</span><span class="s1">&#39;extxyz&#39;</span><span class="p">)</span>
+    <span class="n">n_frames</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">wannier_atoms</span><span class="p">)</span>
+    <span class="n">n_atoms</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">logical_not</span><span class="p">(</span><span class="n">wannier_atoms</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">symbols</span> <span class="o">==</span> <span class="n">wacent_symbol</span><span class="p">))</span>
 
     <span class="n">lumped_dict_c</span> <span class="o">=</span> <span class="n">lumped_dict</span><span class="o">.</span><span class="n">copy</span><span class="p">()</span>
     <span class="n">del_list</span> <span class="o">=</span> <span class="p">[]</span>
@@ -622,7 +629,9 @@ <h1>Source code for ai2_kit.feat.spectrum.viber</h1><div class="highlight"><pre>
     <span class="n">wfc_compute_polar</span> <span class="o">=</span> <span class="n">_set_lumped_wfc</span><span class="p">(</span><span class="n">stc_list</span><span class="p">,</span> <span class="n">lumped_dict_c</span><span class="p">,</span> <span class="n">cutoff</span><span class="p">,</span> <span class="n">wacent_symbol</span><span class="p">,</span> <span class="n">to_polar</span> <span class="o">=</span> <span class="kc">True</span><span class="p">)</span>
     <span class="n">wfc_save</span> <span class="o">=</span> <span class="n">_set_lumped_wfc</span><span class="p">(</span><span class="n">stc_list</span><span class="p">,</span> <span class="n">lumped_dict_c</span><span class="p">,</span> <span class="n">cutoff</span><span class="p">,</span> <span class="n">wacent_symbol</span><span class="p">,</span> <span class="n">to_polar</span> <span class="o">=</span> <span class="kc">False</span><span class="p">)</span>
 
-    <span class="n">dp_sys</span><span class="o">.</span><span class="n">data</span><span class="p">[</span><span class="s1">&#39;atomic_dipole&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">wfc_save</span>
+    <span class="n">wannier</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">([</span><span class="n">n_frames</span><span class="p">,</span> <span class="n">n_atoms</span><span class="p">,</span> <span class="mi">3</span><span class="p">])</span>
+    <span class="n">wannier</span><span class="p">[:,</span> <span class="n">sel_ids</span><span class="p">,</span> <span class="p">:]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">wfc_save</span><span class="p">,</span> <span class="p">[</span><span class="n">n_frames</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">])</span>
+    <span class="n">dp_sys</span><span class="o">.</span><span class="n">data</span><span class="p">[</span><span class="s2">&quot;atomic_dipole&quot;</span><span class="p">]</span> <span class="o">=</span> <span class="n">wannier</span>
 
     <span class="k">if</span> <span class="n">mode</span> <span class="o">==</span> <span class="s1">&#39;both&#39;</span><span class="p">:</span>
         <span class="n">wannier_atoms_x</span> <span class="o">=</span> <span class="n">ase</span><span class="o">.</span><span class="n">io</span><span class="o">.</span><span class="n">read</span><span class="p">(</span><span class="n">os</span><span class="o">.</span><span class="n">path</span><span class="o">.</span><span class="n">join</span><span class="p">(</span><span class="n">data_dir</span><span class="p">,</span> <span class="n">wannier_x</span><span class="p">),</span> <span class="n">index</span><span class="o">=</span><span class="s2">&quot;:&quot;</span><span class="p">,</span> <span class="nb">format</span><span class="o">=</span><span class="s1">&#39;extxyz&#39;</span><span class="p">)</span>
@@ -643,9 +652,9 @@ <h1>Source code for ai2_kit.feat.spectrum.viber</h1><div class="highlight"><pre>
         <span class="n">polar</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">wfc_y</span> <span class="o">-</span> <span class="n">wfc_compute_polar</span><span class="p">)</span> <span class="o">/</span> <span class="n">eps</span>
         <span class="n">polar</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">2</span><span class="p">]</span> <span class="o">=</span> <span class="p">(</span><span class="n">wfc_z</span> <span class="o">-</span> <span class="n">wfc_compute_polar</span><span class="p">)</span> <span class="o">/</span> <span class="n">eps</span>
 
-        <span class="n">dp_sys</span><span class="o">.</span><span class="n">data</span><span class="p">[</span><span class="s1">&#39;atomic_polarizability&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">polar</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">polar</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span>
+        <span class="n">dp_sys</span><span class="o">.</span><span class="n">data</span><span class="p">[</span><span class="s1">&#39;atomic_polarizability&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">polar</span><span class="p">,</span> <span class="p">[</span><span class="n">n_frames</span><span class="p">,</span> <span class="n">n_atoms</span><span class="p">,</span> <span class="mi">9</span><span class="p">])</span>
     <span class="k">elif</span> <span class="n">mode</span> <span class="o">==</span> <span class="s1">&#39;dipole_only&#39;</span><span class="p">:</span>
-        <span class="n">dp_sys</span><span class="o">.</span><span class="n">data</span><span class="p">[</span><span class="s1">&#39;atomic_polarizability&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span><span class="o">-</span><span class="mi">1</span><span class="p">])</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span><span class="o">-</span><span class="mi">1</span><span class="p">)</span>
+        <span class="n">dp_sys</span><span class="o">.</span><span class="n">data</span><span class="p">[</span><span class="s1">&#39;atomic_polarizability&#39;</span><span class="p">]</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">full</span><span class="p">((</span><span class="n">n_frames</span><span class="p">,</span> <span class="n">n_atoms</span><span class="p">,</span> <span class="mi">9</span><span class="p">),</span> <span class="o">-</span><span class="mf">1.0</span><span class="p">)</span>
     <span class="k">else</span><span class="p">:</span>
         <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;There is no mode called &#39;</span><span class="si">{</span><span class="n">mode</span><span class="si">}</span><span class="s2">&#39;, expected &#39;both&#39; or &#39;dipole_only&#39;&quot;</span><span class="p">)</span>