Merge pull request #973 from psavery/material-structure-charge

Add ability to modify a material's charge

hexrd/ui/material_site_editor.py

@@ -3,8 +3,11 @@
 import numpy as np
 
 from PySide2.QtCore import QItemSelectionModel, QObject, Signal
-from PySide2.QtWidgets import QLineEdit, QSizePolicy, QTableWidgetItem
+from PySide2.QtWidgets import (
+    QComboBox, QLineEdit, QSizePolicy, QTableWidgetItem
+)
 
+from hexrd.constants import chargestate
 from hexrd.material import Material
 
 from hexrd.ui.periodic_table_dialog import PeriodicTableDialog
@@ -16,10 +19,12 @@
 
 COLUMNS = {
     'symbol': 0,
-    'occupancy': 1,
-    'thermal_factor': 2
+    'charge': 1,
+    'occupancy': 2,
+    'thermal_factor': 3
 }
 
+DEFAULT_CHARGE = '0'
 DEFAULT_U = Material.DFLT_U[0]
 
 OCCUPATION_MIN = 0
@@ -44,6 +49,7 @@ def __init__(self, site, parent=None):
 
         self._site = site
 
+        self.charge_comboboxes = []
         self.occupancy_spinboxes = []
         self.thermal_factor_spinboxes = []
 
@@ -143,12 +149,17 @@ def atom_types(self, v):
         # Reset all the occupancies
         atoms = self.atoms
         previous_u_values = {x['symbol']: x['U'] for x in atoms}
+        previous_charges = {x['symbol']: x['charge'] for x in atoms}
         atoms.clear()
 
         for symbol in v:
-            # Use the previous U if available. Otherwise, use the default.
-            U = previous_u_values.get(symbol, DEFAULT_U)
-            atoms.append({'symbol': symbol, 'U': U})
+            # Use previous values if available. Otherwise, use the defaults.
+            atom = {
+                'symbol': symbol,
+                'U': previous_u_values.get(symbol, DEFAULT_U),
+                'charge': previous_charges.get(symbol, DEFAULT_CHARGE),
+            }
+            atoms.append(atom)
 
         self.reset_occupancies()
         self.update_table()
@@ -195,6 +206,22 @@ def create_symbol_label(self, v):
         w = QTableWidgetItem(v)
         return w
 
+    def create_charge_combobox(self, charge, symbol):
+        cb = QComboBox(self.ui.table)
+
+        if charge not in chargestate[symbol]:
+            raise Exception(f'Invalid charge {charge} for {symbol}')
+
+        cb.addItems(chargestate[symbol])
+        cb.setCurrentText(charge)
+        cb.currentIndexChanged.connect(self.update_config)
+
+        size_policy = QSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
+        cb.setSizePolicy(size_policy)
+
+        self.charge_comboboxes.append(cb)
+        return cb
+
     def create_occupancy_spinbox(self, v):
         sb = ScientificDoubleSpinBox(self.ui.table)
         sb.setKeyboardTracking(False)
@@ -225,6 +252,7 @@ def create_thermal_factor_spinbox(self, v):
         return sb
 
     def clear_table(self):
+        self.charge_comboboxes.clear()
         self.occupancy_spinboxes.clear()
         self.thermal_factor_spinboxes.clear()
         self.ui.table.clearContents()
@@ -260,6 +288,9 @@ def update_table(self):
                 w = self.create_symbol_label(atom['symbol'])
                 self.ui.table.setItem(i, COLUMNS['symbol'], w)
 
+                w = self.create_charge_combobox(atom['charge'], atom['symbol'])
+                self.ui.table.setCellWidget(i, COLUMNS['charge'], w)
+
                 w = self.create_occupancy_spinbox(atom['occupancy'])
                 self.ui.table.setCellWidget(i, COLUMNS['occupancy'], w)
 
@@ -293,6 +324,9 @@ def update_config(self):
         for i, w in enumerate(self.fractional_coords_widgets):
             self.fractional_coords[i] = w.value()
 
+        for atom, combobox in zip(self.atoms, self.charge_comboboxes):
+            atom['charge'] = combobox.currentText()
+
         for atom, spinbox in zip(self.atoms, self.occupancy_spinboxes):
             atom['occupancy'] = spinbox.value()
 

hexrd/ui/material_structure_editor.py

@@ -23,11 +23,13 @@
     'atoms': [
         {
             'symbol': 'O',
+            'charge': '0',
             'occupancy': 0.5,
             'U': 1e-6
         },
         {
             'symbol': 'Ce',
+            'charge': '0',
             'occupancy': 0.5,
             'U': 1e-6
         }
@@ -234,13 +236,15 @@ def update_material(self):
         # Convert the sites back to the material data format
         info_array = []
         type_array = []
+        charge_array = []
         U_array = []
 
         for site in self.sites:
             for atom in site['atoms']:
                 info_array.append((*site['fractional_coords'],
                                    atom['occupancy']))
                 type_array.append(ptable[atom['symbol']])
+                charge_array.append(atom['charge'])
                 U_array.append(atom['U'])
 
         if any(isinstance(x, np.ndarray) for x in U_array):
@@ -249,7 +253,7 @@ def update_material(self):
                 U_array[i] = scalar_to_tensor(U)
 
         mat = self.material
-        mat._set_atomdata(type_array, info_array, U_array)
+        mat._set_atomdata(type_array, info_array, U_array, charge_array)
 
         self.material_modified.emit()
 
@@ -263,11 +267,13 @@ def generate_sites(self):
             'atoms': [
                 {
                     'symbol': 'Na',
+                    'charge': '0',
                     'occupancy': 0.5,
                     'U': 4.18e-7
                 },
                 {
                     'symbol': 'Br',
+                    'charge': '0',
                     'occupancy': 0.5,
                     'U': 4.18e-7
                 }
@@ -287,6 +293,7 @@ def coords_equal(v1, v2, tol=1.e-5):
 
         info_array = mat._atominfo
         type_array = mat._atomtype
+        charge_array = mat.charge
         U_array = mat._U
 
         if any(isinstance(x, np.ndarray) for x in U_array):
@@ -322,6 +329,7 @@ def coords_equal(v1, v2, tol=1.e-5):
             for i in indices:
                 atom = {
                     'symbol': ptableinverse[type_array[i]],
+                    'charge': charge_array[i],
                     'occupancy': info_array[i][3],
                     'U': U_array[i]
                 }

hexrd/ui/resources/ui/material_site_editor.ui

@@ -52,6 +52,11 @@
           <string>Symbol</string>
          </property>
         </column>
+        <column>
+         <property name="text">
+          <string>Charge</string>
+         </property>
+        </column>
         <column>
          <property name="text">
           <string>Occupancy</string>