@@ -750,7 +750,118 @@ def get_dims_with_index_levels(
750750 return dims_with_levels
751751
752752
753- def get_label_position (
753+ class LabelPositionIndex :
754+ """
755+ Index for fast O(log n) lookup of label positions using binary search.
756+
757+ This class builds a sorted index of label ranges and uses binary search
758+ to find which container (variable/constraint) a label belongs to.
759+
760+ Parameters
761+ ----------
762+ obj : Any
763+ Container object with items() method returning (name, val) pairs,
764+ where val has .labels and .range attributes.
765+ """
766+
767+ __slots__ = ("_starts" , "_names" , "_obj" , "_built" )
768+
769+ def __init__ (self , obj : Any ) -> None :
770+ self ._obj = obj
771+ self ._starts : np .ndarray | None = None
772+ self ._names : list [str ] | None = None
773+ self ._built = False
774+
775+ def _build_index (self ) -> None :
776+ """Build the sorted index of label ranges."""
777+ if self ._built :
778+ return
779+
780+ ranges = []
781+ for name , val in self ._obj .items ():
782+ start , stop = val .range
783+ ranges .append ((start , name ))
784+
785+ # Sort by start value
786+ ranges .sort (key = lambda x : x [0 ])
787+ self ._starts = np .array ([r [0 ] for r in ranges ])
788+ self ._names = [r [1 ] for r in ranges ]
789+ self ._built = True
790+
791+ def invalidate (self ) -> None :
792+ """Invalidate the index (call when items are added/removed)."""
793+ self ._built = False
794+ self ._starts = None
795+ self ._names = None
796+
797+ def find_single (self , value : int ) -> tuple [str , dict ] | tuple [None , None ]:
798+ """Find the name and coordinates for a single label value."""
799+ if value == - 1 :
800+ return None , None
801+
802+ self ._build_index ()
803+ starts = self ._starts
804+ names = self ._names
805+ assert starts is not None and names is not None
806+
807+ # Binary search to find the right range
808+ idx = int (np .searchsorted (starts , value , side = "right" )) - 1
809+
810+ if idx < 0 or idx >= len (starts ):
811+ raise ValueError (f"Label { value } )
812+
813+ name = names [idx ]
814+ val = self ._obj [name ]
815+ start , stop = val .range
816+
817+ # Verify the value is in range
818+ if value < start or value >= stop :
819+ raise ValueError (f"Label { value } )
820+
821+ labels = val .labels
822+ index = np .unravel_index (value - start , labels .shape )
823+ coord = {dim : labels .indexes [dim ][i ] for dim , i in zip (labels .dims , index )}
824+ return name , coord
825+
826+ def find_single_with_index (
827+ self , value : int
828+ ) -> tuple [str , dict , tuple [int , ...]] | tuple [None , None , None ]:
829+ """
830+ Find name, coordinates, and raw numpy index for a single label value.
831+
832+ Returns (name, coord, index) where index is a tuple of integers that
833+ can be used for direct numpy indexing (e.g., arr.values[index]).
834+ This avoids the overhead of xarray's .sel() method.
835+ """
836+ if value == - 1 :
837+ return None , None , None
838+
839+ self ._build_index ()
840+ starts = self ._starts
841+ names = self ._names
842+ assert starts is not None and names is not None
843+
844+ # Binary search to find the right range
845+ idx = int (np .searchsorted (starts , value , side = "right" )) - 1
846+
847+ if idx < 0 or idx >= len (starts ):
848+ raise ValueError (f"Label { value } )
849+
850+ name = names [idx ]
851+ val = self ._obj [name ]
852+ start , stop = val .range
853+
854+ # Verify the value is in range
855+ if value < start or value >= stop :
856+ raise ValueError (f"Label { value } )
857+
858+ labels = val .labels
859+ index = np .unravel_index (value - start , labels .shape )
860+ coord = {dim : labels .indexes [dim ][i ] for dim , i in zip (labels .dims , index )}
861+ return name , coord , index
862+
863+
864+ def _get_label_position_linear (
754865 obj : Any , values : int | np .ndarray
755866) -> (
756867 tuple [str , dict ]
@@ -760,6 +871,9 @@ def get_label_position(
760871):
761872 """
762873 Get tuple of name and coordinate for variable labels.
874+
875+ This is the original O(n) implementation that scans through all items.
876+ Used only for testing/benchmarking comparisons.
763877 """
764878
765879 def find_single (value : int ) -> tuple [str , dict ] | tuple [None , None ]:
@@ -795,6 +909,53 @@ def find_single(value: int) -> tuple[str, dict] | tuple[None, None]:
795909 raise ValueError ("Array's with more than two dimensions is not supported" )
796910
797911
912+ def get_label_position (
913+ obj : Any ,
914+ values : int | np .ndarray ,
915+ index : LabelPositionIndex | None = None ,
916+ ) -> (
917+ tuple [str , dict ]
918+ | tuple [None , None ]
919+ | list [tuple [str , dict ] | tuple [None , None ]]
920+ | list [list [tuple [str , dict ] | tuple [None , None ]]]
921+ ):
922+ """
923+ Get tuple of name and coordinate for variable labels.
924+
925+ Uses O(log n) binary search with a cached index for fast lookups.
926+
927+ Parameters
928+ ----------
929+ obj : Any
930+ Container object with items() method (Variables or Constraints).
931+ values : int or np.ndarray
932+ Label value(s) to look up.
933+ index : LabelPositionIndex, optional
934+ Pre-built index for fast lookups. If None, one will be created.
935+
936+ Returns
937+ -------
938+ tuple or list
939+ (name, coord) tuple for single values, or list of tuples for arrays.
940+ """
941+ if index is None :
942+ index = LabelPositionIndex (obj )
943+
944+ if isinstance (values , int ):
945+ return index .find_single (values )
946+
947+ values = np .array (values )
948+ ndim = values .ndim
949+ if ndim == 0 :
950+ return index .find_single (values .item ())
951+ elif ndim == 1 :
952+ return [index .find_single (int (v )) for v in values ]
953+ elif ndim == 2 :
954+ return [[index .find_single (int (v )) for v in col ] for col in values .T ]
955+ else :
956+ raise ValueError ("Array's with more than two dimensions is not supported" )
957+
958+
798959def print_coord (coord : dict [str , Any ] | Iterable [Any ]) -> str :
799960 """
800961 Format coordinates into a string representation.
@@ -838,14 +999,16 @@ def print_single_variable(model: Any, label: int) -> str:
838999 return "None"
8391000
8401001 variables = model .variables
841- name , coord = variables .get_label_position (label )
1002+ name , coord , index = variables .get_label_position_with_index (label )
8421003
843- lower = variables [name ].lower .sel (coord ).item ()
844- upper = variables [name ].upper .sel (coord ).item ()
1004+ var = variables [name ]
1005+ # Use direct numpy indexing instead of .sel() for performance
1006+ lower = var .lower .values [index ]
1007+ upper = var .upper .values [index ]
8451008
846- if variables [ name ] .attrs ["binary" ]:
1009+ if var .attrs ["binary" ]:
8471010 bounds = " ∈ {0, 1}"
848- elif variables [ name ] .attrs ["integer" ]:
1011+ elif var .attrs ["integer" ]:
8491012 bounds = f" ∈ Z ⋂ [{ lower :.4g} { upper :.4g}
8501013 else :
8511014 bounds = f" ∈ [{ lower :.4g} { upper :.4g}
0 commit comments