33from collections .abc import Iterable , Mapping , Sequence
44from functools import lru_cache , partial
55from itertools import chain , product , zip_longest
6- from typing import Optional , Union
6+ from typing import Optional , TypeVar , Union
77from warnings import warn
88
99import ixmp
1010import numpy as np
1111import pandas as pd
1212from ixmp .backend import ItemType
13+ from ixmp .backend .jdbc import JDBCBackend
1314from ixmp .util import as_str_list , maybe_check_out , maybe_commit
1415
16+ from message_ix .util .ixmp4 import on_ixmp4backend
17+
18+ # from message_ix.util.scenario_data import PARAMETERS
19+
1520log = logging .getLogger (__name__ )
1621
1722# Also print warnings to stderr
@@ -62,7 +67,7 @@ def __init__(
6267 # Utility methods used by .equ(), .par(), .set(), and .var()
6368
6469 @lru_cache ()
65- def _year_idx (self , name ):
70+ def _year_idx (self , name : str ):
6671 """Return a sequence of (idx_set, idx_name) for 'year'-indexed dims.
6772
6873 Since item dimensionality does not change, the the return value is
@@ -78,7 +83,10 @@ def _year_idx(self, name):
7883 )
7984 )
8085
81- def _year_as_int (self , name , df ):
86+ data_type = TypeVar ("data_type" , pd .Series , pd .DataFrame , dict )
87+
88+ # NOTE super().equ() etc hint that pd objects return, but they may also return dict
89+ def _year_as_int (self , name : str , data : data_type ) -> data_type :
8290 """Convert 'year'-indexed columns of *df* to :obj:`int` dtypes.
8391
8492 :meth:`_year_idx` is used to retrieve a sequence of (idx_set, idx_name)
@@ -90,12 +98,18 @@ def _year_as_int(self, name, df):
9098 year_idx = self ._year_idx (name )
9199
92100 if len (year_idx ):
93- return df .astype ({col_name : "int" for _ , col_name in year_idx })
101+ assert isinstance (data , pd .DataFrame )
102+ # NOTE With the IXMP4Backend, we call scenario.par() for parameters that
103+ # might be empty, which fails df.astype() below.
104+ if data .empty :
105+ return data
106+ return data .astype ({col_name : "int" for _ , col_name in year_idx })
94107 elif name == "year" :
95108 # The 'year' set itself
96- return df .astype (int )
109+ assert isinstance (data , pd .Series )
110+ return data .astype (int )
97111 else :
98- return df
112+ return data
99113
100114 # Override ixmp methods to convert 'year'-indexed columns to int
101115
@@ -251,6 +265,28 @@ def add_par(
251265 # accepts int for "year"-like dimensions. Proxy the call to avoid type check
252266 # failures.
253267 # TODO Move this upstream, to ixmp
268+
269+ if on_ixmp4backend (self ):
270+ from message_ix .util .scenario_setup import check_existence_of_units
271+
272+ # Check for existence of required units
273+ # NOTE these checks are similar to those in super().add_par(), but we
274+ # need them here for access to self
275+ # NOTE it seems to me that only dict or pd.DataFrame key_or_data could
276+ # contain 'unit' already, else it's supplied by keyword or defaults to
277+ # "???"
278+ if isinstance (key_or_data , dict ):
279+ _data = pd .DataFrame .from_dict (key_or_data , orient = "columns" )
280+ elif isinstance (key_or_data , pd .DataFrame ):
281+ _data = key_or_data
282+ else :
283+ _data = pd .DataFrame ()
284+
285+ if "unit" not in _data .columns :
286+ _data ["unit" ] = unit or "???"
287+
288+ check_existence_of_units (platform = self .platform , data = _data )
289+
254290 super ().add_par (name , key_or_data , value , unit , comment ) # type: ignore [arg-type]
255291
256292 add_par .__doc__ = ixmp .Scenario .add_par .__doc__
@@ -323,6 +359,7 @@ def recurse(k, v, parent="World"):
323359 recurse (k , v )
324360
325361 self .add_set ("node" , nodes )
362+ # TODO do we handle levels being added multiple times correctly for ixmp4?
326363 self .add_set ("lvl_spatial" , levels )
327364 self .add_set ("map_spatial_hierarchy" , hierarchy )
328365
@@ -420,8 +457,11 @@ def add_horizon(
420457 # Add the year set elements and first model year
421458 year = sorted (year )
422459 self .add_set ("year" , year )
460+
461+ # Avoid removing default data on IXMP4Backend
462+ is_unique = True if isinstance (self .platform ._backend , JDBCBackend ) else False
423463 self .add_cat (
424- "year" , "firstmodelyear" , firstmodelyear or year [0 ], is_unique = True
464+ "year" , "firstmodelyear" , firstmodelyear or year [0 ], is_unique = is_unique
425465 )
426466
427467 # Calculate the duration of all periods
@@ -560,7 +600,7 @@ def vintage_and_active_years(
560600 vintages = sorted (
561601 self .par (
562602 "technical_lifetime" ,
563- filters = {"node_loc" : ya_args [0 ], "technology" : ya_args [1 ]},
603+ filters = {"node_loc" : [ ya_args [0 ]] , "technology" : [ ya_args [1 ] ]},
564604 )["year_vtg" ].unique ()
565605 )
566606 ya_max = max (vintages ) if tl_only else np .inf
@@ -585,6 +625,7 @@ def vintage_and_active_years(
585625
586626 # Minimum value for year_act
587627 if "in_horizon" in kwargs :
628+ # FIXME I don't receive this in the tutorials
588629 warn (
589630 "'in_horizon' argument to .vintage_and_active_years() will be removed "
590631 "in message_ix>=4.0. Use .query(…) instead per documentation examples." ,
@@ -842,3 +883,16 @@ def rename(self, name: str, mapping: Mapping[str, str], keep: bool = False) -> N
842883 self .remove_set (name , list (mapping .keys ()))
843884
844885 maybe_commit (self , commit , f"Rename { name !r} { mapping } )
886+
887+ def commit (self , comment : str ) -> None :
888+ from message_ix .util .scenario_setup import compose_maps
889+
890+ # JDBCBackend calls these functions as part of every commit, but they have moved
891+ # to message_ix because they handle message-specific data
892+
893+ # The sanity checks fail for some tests (e.g. 'node' being empty)
894+ # ensure_required_indexsets_have_data(scenario=self)
895+
896+ compose_maps (self )
897+
898+ return super ().commit (comment )
0 commit comments