support MultiSystems from_dir

Author: felix5572

dpdata/system.py

@@ -950,7 +950,20 @@ def from_file(cls,file_name,fmt):
         multi_systems = cls()
         multi_systems.load_systems_from_file(file_name=file_name,fmt=fmt)
         return multi_systems
-    
+
+    @classmethod
+    def from_dir(cls,dir_name, file_name, fmt='auto'):
+        multi_systems = cls()
+        target_file_list = []
+        for ii in os.walk(dir_name):
+            if file_name in ii[2]:
+                target_file = os.path.join(ii[0], file_name)
+                target_file_list.append(target_file)
+        for target_file in target_file_list:
+            multi_systems.append(LabeledSystem(file_name=target_file, fmt=fmt))
+        return multi_systems
+
+
     def load_systems_from_file(self, file_name=None, fmt=None):
         if file_name is not None:
             if fmt is None:
@@ -1015,10 +1028,10 @@ def check_atom_names(self, system):
             system.add_atom_names(new_in_self)
         system.sort_atom_names()
 
-    def from_quip_gap_xyz_file(self,filename):
-        # quip_gap_xyz_systems = QuipGapxyzSystems(filename)
+    def from_quip_gap_xyz_file(self,file_name):
+        # quip_gap_xyz_systems = QuipGapxyzSystems(file_name)
         # print(next(quip_gap_xyz_systems))
-        for info_dict in QuipGapxyzSystems(filename):
+        for info_dict in QuipGapxyzSystems(file_name):
             system=LabeledSystem(data=info_dict)
             self.append(system)
 

dpdata/xyz/quip_gap_xyz.py

@@ -1,9 +1,4 @@
 #!/usr/bin/env python3
-
-#%%
-# with open('./test.xyz', 'r') as xyz_file:
-#     lines = xyz_file.readlines()
-#     print(lines)
 #%% 
 import numpy as np
 from collections import OrderedDict
@@ -44,20 +39,16 @@ def handle_single_xyz_frame(lines):
         if len(lines) != atom_num + 2:
             raise RuntimeError("format error, atom_num=={}, {}!=atom_num+2".format(atom_num, len(lines)))
         data_format_line = lines[1].strip('\n').strip()+str(' ')
-        p1 = re.compile(r'(?P<key>\S+)=(?P<quote>[\'\"]?)(?P<value>.*?)(?P=quote)\s+')
-        p2 = re.compile(r'(?P<key>\w+?):(?P<datatype>[a-zA-Z]):(?P<value>\d+)')
-        field_list = [kv_dict.groupdict() for kv_dict in p1.finditer(data_format_line)]
-        field_dict = {}
-        for item in field_list:
-            field_dict[item['key']]=item['value']
-        data_format_line = lines[1]
-        data_format_list= [m.groupdict() for m in p1.finditer(data_format_line)]
+        field_value_pattern= re.compile(r'(?P<key>\S+)=(?P<quote>[\'\"]?)(?P<value>.*?)(?P=quote)\s+')
+        prop_pattern = re.compile(r'(?P<key>\w+?):(?P<datatype>[a-zA-Z]):(?P<value>\d+)')
+
+        data_format_list= [kv_dict.groupdict() for kv_dict in field_value_pattern.finditer(data_format_line)]
         field_dict = {}
         for item in data_format_list:
             field_dict[item['key']]=item['value']
 
         Properties = field_dict['Properties']
-        prop_list = [m.groupdict() for m in p2.finditer(Properties)]
+        prop_list = [kv_dict.groupdict() for kv_dict in prop_pattern.finditer(Properties)]
 
         data_lines = []
         for line in lines[2:]:
@@ -127,6 +118,7 @@ def handle_single_xyz_frame(lines):
             virials = np.array([np.array(list(filter(bool,field_dict['virial'].split(' ')))).reshape(3,3)]).astype('float32')
         else:
             virials = None
+
         info_dict = {}
         info_dict['atom_names'] = list(type_num_array[:,0])
         info_dict['atom_numbs'] = list(type_num_array[:,1].astype(int))