1515import numpy as np
1616from numpy .lib .recfunctions import stack_arrays
1717
18- from .utils import MfList , OptionBlock , Transient2d , Util2d , Util3d , check
18+ from .utils import (
19+ MfList ,
20+ OptionBlock ,
21+ Transient2d ,
22+ Util2d ,
23+ Util3d ,
24+ check ,
25+ create_empty_recarray ,
26+ )
1927from .utils .flopy_io import ulstrd
2028
2129
@@ -1085,6 +1093,9 @@ def load(
10851093 bnd_output_cln = None
10861094 stress_period_data_cln = {}
10871095 current_cln = None
1096+ is_mfusgwel = "mfusgwel" in pak_type_str
1097+ max_comment_cols = 0
1098+ max_comment_cols_cln = 0
10881099 for iper in range (nper ):
10891100 if model .verbose :
10901101 msg = f" loading { pak_type } for kper { iper + 1 :5d} "
@@ -1126,7 +1137,36 @@ def load(
11261137 current = pak_type .get_empty (
11271138 itmp , aux_names = aux_names , structured = model .structured , ** usg_args
11281139 )
1129- current = ulstrd (f , itmp , current , model , sfac_columns , ext_unit_dict )
1140+ if is_mfusgwel :
1141+ current , extra_tokens = ulstrd (
1142+ f , itmp , current , model , sfac_columns , ext_unit_dict ,
1143+ capture_comments = True ,
1144+ )
1145+ n_comments = max ((len (ct ) for ct in extra_tokens ), default = 0 )
1146+ if n_comments > 0 :
1147+ max_comment_cols = max (max_comment_cols , n_comments )
1148+ new_dtype = current .dtype
1149+ for ci in range (n_comments ):
1150+ new_dtype = Package .add_to_dtype (
1151+ new_dtype , f"comment{ ci + 1 } " , object
1152+ )
1153+ new_ra = create_empty_recarray (
1154+ len (current ), new_dtype , default_value = - 1.0e10
1155+ )
1156+ for name in current .dtype .names :
1157+ new_ra [name ] = current [name ]
1158+ for ii in range (len (current )):
1159+ for ci in range (n_comments ):
1160+ cname = f"comment{ ci + 1 } "
1161+ if ci < len (extra_tokens [ii ]):
1162+ new_ra [cname ][ii ] = extra_tokens [ii ][ci ]
1163+ else :
1164+ new_ra [cname ][ii ] = ""
1165+ current = new_ra
1166+ else :
1167+ current = ulstrd (
1168+ f , itmp , current , model , sfac_columns , ext_unit_dict
1169+ )
11301170 if model .structured :
11311171 current ["k" ] -= 1
11321172 current ["i" ] -= 1
@@ -1149,9 +1189,41 @@ def load(
11491189 current_cln = pak_type .get_empty (
11501190 itmp_cln , aux_names = aux_names , structured = False , ** usg_args
11511191 )
1152- current_cln = ulstrd (
1153- f , itmp_cln , current_cln , model , sfac_columns , ext_unit_dict
1154- )
1192+ if is_mfusgwel :
1193+ current_cln , extra_tokens_cln = ulstrd (
1194+ f , itmp_cln , current_cln , model , sfac_columns ,
1195+ ext_unit_dict , capture_comments = True ,
1196+ )
1197+ n_comments_cln = max (
1198+ (len (ct ) for ct in extra_tokens_cln ), default = 0
1199+ )
1200+ if n_comments_cln > 0 :
1201+ max_comment_cols_cln = max (
1202+ max_comment_cols_cln , n_comments_cln
1203+ )
1204+ new_dtype = current_cln .dtype
1205+ for ci in range (n_comments_cln ):
1206+ new_dtype = Package .add_to_dtype (
1207+ new_dtype , f"comment{ ci + 1 } " , object
1208+ )
1209+ new_ra = create_empty_recarray (
1210+ len (current_cln ), new_dtype , default_value = - 1.0e10
1211+ )
1212+ for name in current_cln .dtype .names :
1213+ new_ra [name ] = current_cln [name ]
1214+ for ii in range (len (current_cln )):
1215+ for ci in range (n_comments_cln ):
1216+ cname = f"comment{ ci + 1 } "
1217+ if ci < len (extra_tokens_cln [ii ]):
1218+ new_ra [cname ][ii ] = extra_tokens_cln [ii ][ci ]
1219+ else :
1220+ new_ra [cname ][ii ] = ""
1221+ current_cln = new_ra
1222+ else :
1223+ current_cln = ulstrd (
1224+ f , itmp_cln , current_cln , model , sfac_columns ,
1225+ ext_unit_dict ,
1226+ )
11551227 current_cln ["node" ] -= 1
11561228 bnd_output_cln = np .recarray .copy (current_cln )
11571229 else :
@@ -1231,6 +1303,45 @@ def load(
12311303 0 , aux_names = aux_names , structured = model .structured , ** usg_args
12321304 ).dtype
12331305
1306+ # Normalize comment columns across stress periods for MfUsgWel
1307+ if is_mfusgwel and (max_comment_cols > 0 or max_comment_cols_cln > 0 ):
1308+ for spd_dict , max_cc in [
1309+ (stress_period_data , max_comment_cols ),
1310+ (stress_period_data_cln , max_comment_cols_cln ),
1311+ ]:
1312+ if max_cc == 0 :
1313+ continue
1314+ for iper in spd_dict :
1315+ spd = spd_dict [iper ]
1316+ if not isinstance (spd , np .recarray ):
1317+ continue
1318+ existing = sum (
1319+ 1 for n in spd .dtype .names if n .startswith ("comment" )
1320+ )
1321+ if existing < max_cc :
1322+ new_dtype = spd .dtype
1323+ for ci in range (existing , max_cc ):
1324+ new_dtype = Package .add_to_dtype (
1325+ new_dtype , f"comment{ ci + 1 } " , object
1326+ )
1327+ new_ra = create_empty_recarray (
1328+ len (spd ), new_dtype , default_value = - 1.0e10
1329+ )
1330+ for name in spd .dtype .names :
1331+ new_ra [name ] = spd [name ]
1332+ for ci in range (existing , max_cc ):
1333+ cname = f"comment{ ci + 1 } "
1334+ for ii in range (len (new_ra )):
1335+ new_ra [cname ][ii ] = ""
1336+ spd_dict [iper ] = new_ra
1337+
1338+ # Update dtype to include comment columns
1339+ if max_comment_cols > 0 :
1340+ for ci in range (max_comment_cols ):
1341+ dtype = Package .add_to_dtype (
1342+ dtype , f"comment{ ci + 1 } " , object
1343+ )
1344+
12341345 if openfile :
12351346 f .close ()
12361347
@@ -1248,6 +1359,12 @@ def load(
12481359 cln_dtype = pak_type .get_empty (
12491360 0 , aux_names = aux_names , structured = False , ** usg_args
12501361 ).dtype
1362+ # Update cln_dtype to include comment columns
1363+ if max_comment_cols_cln > 0 :
1364+ for ci in range (max_comment_cols_cln ):
1365+ cln_dtype = Package .add_to_dtype (
1366+ cln_dtype , f"comment{ ci + 1 } " , object
1367+ )
12511368 pak = pak_type (
12521369 model ,
12531370 ipakcb = ipakcb ,
0 commit comments