Chloe lee doper

smact/dopant_prediction/doper.py

@@ -1,6 +1,6 @@
 import os
 from itertools import groupby
-from typing import List, Optional, Tuple, Type
+from typing import List, Optional, Tuple
 
 import numpy as np
 from pymatgen.util import plotting
@@ -289,11 +289,6 @@ def get_dopants(
                 groupby_lists[i] = grouped_top_data
                 del grouped_data
 
-        # select top n elements
-        dopants_lists = [
-            dopants_list[:num_dopants] for dopants_list in dopants_lists
-        ]
-
         keys = [
             "n-type cation substitutions",
             "p-type cation substitutions",
@@ -304,7 +299,10 @@ def get_dopants(
         self.results = self._merge_dicts(keys, dopants_lists, groupby_lists)
 
         # return the top (num_dopants) results for each case
-        return self.results
+        return {
+            key: value.get("sorted")[:num_dopants]
+            for key, value in self.results.items()
+        }
 
     def plot_dopants(self) -> None:
         """
@@ -332,8 +330,12 @@ def plot_dopants(self) -> None:
                 }
             plotting.periodic_table_heatmap(
                 elemental_data=dict_results,
-                cmap="rainbow",
-                blank_color="gainsboro",
+                cmap="Reds",
+                cmap_range=(
+                    min(dict_results.values()),
+                    max(dict_results.values()),
+                ),
+                blank_color="#D4D4D4",
                 edge_color="white",
             )
 

smact/tests/test_doper.py

@@ -1,17 +1,15 @@
 import os
 import unittest
 
-import smact
 from smact.dopant_prediction import doper
-from smact.structure_prediction import mutation, utilities
+from smact.structure_prediction import utilities
 
 files_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "files")
 TEST_LAMBDA_JSON = os.path.join(files_dir, "test_lambda_tab.json")
 
 
 class DopantPredictionTest(unittest.TestCase):
     def test_dopant_prediction(self):
-        num_dopants = 10
         test_specie = ("Cu+", "Ga3+", "S2-")
         test = doper.Doper(test_specie)
 
@@ -25,7 +23,8 @@ def test_dopant_prediction(self):
         _, an_charge = utilities.parse_spec(anion_min_charge)
 
         # Assert: Length of the list and return type (dictionary: list)
-        result = test.get_dopants()
+        dopants = test.get_dopants()
+        result = test.results
         self.assertIs(type(result), dict)
         for d in result.values():
             self.assertIn("sorted", d)
@@ -59,7 +58,8 @@ def test_dopant_prediction_skipspecies(self):
         test = doper.Doper(
             test_specie, embedding="skipspecies", use_probability=False
         )
-        result = test.get_dopants()
+        dopants = test.get_dopants()
+        result = test.results
 
         n_sub_list_cat = result.get("n-type cation substitutions").get("sorted")
         p_sub_list_cat = result.get("p-type cation substitutions").get("sorted")