job-shop-lib 1.0.0b3__py3-none-any.whl → 1.0.0b5__py3-none-any.whl

job_shop_lib/__init__.py

@@ -19,7 +19,7 @@ from job_shop_lib._schedule import Schedule
 from job_shop_lib._base_solver import BaseSolver, Solver
 
 
-__version__ = "1.0.0-b.3"
+__version__ = "1.0.0-b.5"
 
 __all__ = [
     "Operation",

job_shop_lib/reinforcement_learning/__init__.py

@@ -13,6 +13,8 @@
     RenderConfig
     add_padding
     create_edge_type_dict
+    ResourceTaskGraphObservation
+    ResourceTaskGraphObservationDict
 
 """
 
@@ -31,6 +33,7 @@ from job_shop_lib.reinforcement_learning._reward_observers import (
 from job_shop_lib.reinforcement_learning._utils import (
     add_padding,
     create_edge_type_dict,
+    map_values,
 )
 
 from job_shop_lib.reinforcement_learning._single_job_shop_graph_env import (
@@ -39,6 +42,9 @@ from job_shop_lib.reinforcement_learning._single_job_shop_graph_env import (
 from job_shop_lib.reinforcement_learning._multi_job_shop_graph_env import (
     MultiJobShopGraphEnv,
 )
+from ._resource_task_graph_observation import (
+    ResourceTaskGraphObservation, ResourceTaskGraphObservationDict
+)
 
 
 __all__ = [
@@ -52,4 +58,7 @@ __all__ = [
     "add_padding",
     "MultiJobShopGraphEnv",
     "create_edge_type_dict",
+    "ResourceTaskGraphObservation",
+    "map_values",
+    "ResourceTaskGraphObservationDict",
 ]

job_shop_lib/reinforcement_learning/_multi_job_shop_graph_env.py

@@ -235,13 +235,13 @@ class MultiJobShopGraphEnv(gym.Env):
     @ready_operations_filter.setter
     def ready_operations_filter(
         self,
-        pruning_function: Callable[
+        ready_operations_filter: Callable[
             [Dispatcher, List[Operation]], List[Operation]
         ],
     ) -> None:
         """Sets the ready operations filter."""
         self.single_job_shop_graph_env.dispatcher.ready_operations_filter = (
-            pruning_function
+            ready_operations_filter
         )
 
     @property

job_shop_lib/reinforcement_learning/_resource_task_graph_observation.py

@@ -0,0 +1,251 @@
+"""Contains wrappers for the environments."""
+
+from typing import TypeVar, TypedDict, Generic
+from gymnasium import ObservationWrapper
+import numpy as np
+from numpy.typing import NDArray
+
+from job_shop_lib.reinforcement_learning import (
+    ObservationDict,
+    SingleJobShopGraphEnv,
+    MultiJobShopGraphEnv,
+    create_edge_type_dict,
+    map_values,
+)
+from job_shop_lib.graphs import NodeType
+from job_shop_lib.dispatching.feature_observers import FeatureType
+
+T = TypeVar("T", bound=np.number)
+EnvType = TypeVar(  # pylint: disable=invalid-name
+    "EnvType", bound=SingleJobShopGraphEnv | MultiJobShopGraphEnv
+)
+
+
+class ResourceTaskGraphObservationDict(TypedDict):
+    """Represents a dictionary for resource task graph observations."""
+
+    edge_index_dict: dict[str, NDArray[np.int64]]
+    node_features_dict: dict[str, NDArray[np.float32]]
+    original_ids_dict: dict[str, NDArray[np.int32]]
+
+
+# pylint: disable=line-too-long
+class ResourceTaskGraphObservation(ObservationWrapper, Generic[EnvType]):
+    """Observation wrapper that converts an observation following the
+    :class:`ObservationDict` format to a format suitable to PyG's
+    [`HeteroData`](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.data.HeteroData.html).
+
+    In particular, the ``edge_index`` is converted into a ``edge_index_dict``
+    with keys ``(node_type_i, "to", node_type_j)``. The ``node_type_i`` and
+    ``node_type_j`` are the node types of the source and target nodes,
+    respectively.
+
+    Attributes:
+        global_to_local_id: A dictionary mapping global node IDs to local node
+            IDs for each node type.
+        type_ranges: A dictionary mapping node type names to (start, end) index
+            ranges.
+
+    Args:
+        env: The environment to wrap.
+    """
+
+    def __init__(self, env: EnvType):
+        super().__init__(env)
+        self.env = env  # Unnecessary, but makes mypy happy
+        self.global_to_local_id = self._compute_id_mappings()
+        self.type_ranges = self._compute_node_type_ranges()
+
+    def step(self, action: tuple[int, int]):
+        """Takes a step in the environment.
+
+        Args:
+            action:
+                The action to take. The action is a tuple of two integers
+                (job_id, machine_id):
+                the job ID and the machine ID in which to schedule the
+                operation.
+
+        Returns:
+            A tuple containing the following elements:
+
+            - The observation of the environment.
+            - The reward obtained.
+            - Whether the environment is done.
+            - Whether the episode was truncated (always False).
+            - A dictionary with additional information. The dictionary
+              contains the following keys: "feature_names", the names of the
+              features in the observation; and "available_operations_with_ids",
+              a list of available actions in the form of (operation_id,
+              machine_id, job_id).
+        """
+        observation, reward, done, truncated, info = self.env.step(action)
+        return self.observation(observation), reward, done, truncated, info
+
+    def reset(self, *, seed: int | None = None, options: dict | None = None):
+        """Resets the environment.
+
+        Args:
+            seed:
+                Added to match the signature of the parent class. It is not
+                used in this method.
+            options:
+                Additional options to pass to the environment. Not used in
+                this method.
+
+        Returns:
+            A tuple containing the following elements:
+
+            - The observation of the environment.
+            - A dictionary with additional information, keys
+                include: "feature_names", the names of the features in the
+                observation; and "available_operations_with_ids", a list of
+                available a list of available actions in the form of
+                (operation_id, machine_id, job_id).
+        """
+        observation, info = self.env.reset()
+        return self.observation(observation), info
+
+    def _compute_id_mappings(self) -> dict[int, int]:
+        """Computes mappings from global node IDs to type-local IDs.
+
+        Returns:
+            A dictionary mapping global node IDs to local node IDs for each
+            node type.
+        """
+        mappings = {}
+        for node_type in NodeType:
+            type_nodes = self.unwrapped.job_shop_graph.nodes_by_type[node_type]
+            if not type_nodes:
+                continue
+            # Create mapping from global ID to local ID
+            # (0 to len(type_nodes)-1)
+            type_mapping = {
+                node.node_id: local_id
+                for local_id, node in enumerate(type_nodes)
+            }
+            mappings.update(type_mapping)
+
+        return mappings
+
+    def _compute_node_type_ranges(self) -> dict[str, tuple[int, int]]:
+        """Computes index ranges for each node type.
+
+        Returns:
+            Dictionary mapping node type names to (start, end) index ranges
+        """
+        type_ranges = {}
+        for node_type in NodeType:
+            type_nodes = self.unwrapped.job_shop_graph.nodes_by_type[node_type]
+            if not type_nodes:
+                continue
+            start = min(node.node_id for node in type_nodes)
+            end = max(node.node_id for node in type_nodes) + 1
+            type_ranges[node_type.name.lower()] = (start, end)
+
+        return type_ranges
+
+    def observation(self, observation: ObservationDict):
+        edge_index_dict = create_edge_type_dict(
+            observation["edge_index"],
+            type_ranges=self.type_ranges,
+            relationship="to",
+        )
+        # mapping from global node ID to local node ID
+        for key, edge_index in edge_index_dict.items():
+            edge_index_dict[key] = map_values(
+                edge_index, self.global_to_local_id
+            )
+        node_features_dict = self._create_node_features_dict(observation)
+        node_features_dict, original_ids_dict = self._remove_nodes(
+            node_features_dict, observation["removed_nodes"]
+        )
+
+        return {
+            "edge_index_dict": edge_index_dict,
+            "node_features_dict": node_features_dict,
+            "original_ids_dict": original_ids_dict,
+        }
+
+    def _create_node_features_dict(
+        self, observation: ObservationDict
+    ) -> dict[str, NDArray]:
+        """Creates a dictionary of node features for each node type.
+
+        Args:
+            observation: The observation dictionary.
+
+        Returns:
+            Dictionary mapping node type names to node features.
+        """
+        node_type_to_feature_type = {
+            NodeType.OPERATION: FeatureType.OPERATIONS,
+            NodeType.MACHINE: FeatureType.MACHINES,
+            NodeType.JOB: FeatureType.JOBS,
+        }
+        node_features_dict = {}
+        for node_type, feature_type in node_type_to_feature_type.items():
+            if node_type in self.unwrapped.job_shop_graph.nodes_by_type:
+                node_features_dict[feature_type.value] = observation[
+                    feature_type.value
+                ]
+                continue
+            if feature_type != FeatureType.JOBS:
+                continue
+            assert FeatureType.OPERATIONS.value in observation
+            job_features = observation[
+                feature_type.value  # type: ignore[literal-required]
+            ]
+            job_ids_of_ops = [
+                node.operation.job_id
+                for node in self.unwrapped.job_shop_graph.nodes_by_type[
+                    NodeType.OPERATION
+                ]
+            ]
+            job_features_expanded = job_features[job_ids_of_ops]
+            operation_features = observation[FeatureType.OPERATIONS.value]
+            node_features_dict[FeatureType.OPERATIONS.value] = np.concatenate(
+                (operation_features, job_features_expanded), axis=1
+            )
+        return node_features_dict
+
+    def _remove_nodes(
+        self,
+        node_features_dict: dict[str, NDArray[np.float32]],
+        removed_nodes: NDArray[np.bool_],
+    ) -> tuple[dict[str, NDArray[np.float32]], dict[str, NDArray[np.int32]]]:
+        """Removes nodes from the node features dictionary.
+
+        Args:
+            node_features_dict: The node features dictionary.
+
+        Returns:
+            The node features dictionary with the nodes removed and a
+            dictionary containing the original node ids.
+        """
+        removed_nodes_dict: dict[str, NDArray[np.float32]] = {}
+        original_ids_dict: dict[str, NDArray[np.int32]] = {}
+        feature_type_to_node_type = {
+            FeatureType.OPERATIONS.value: NodeType.OPERATION,
+            FeatureType.MACHINES.value: NodeType.MACHINE,
+            FeatureType.JOBS.value: NodeType.JOB,
+        }
+        for feature_type, features in node_features_dict.items():
+            node_type = feature_type_to_node_type[feature_type].name.lower()
+            if node_type not in self.type_ranges:
+                continue
+            start, end = self.type_ranges[node_type]
+            removed_nodes_of_this_type = removed_nodes[start:end]
+            removed_nodes_dict[node_type] = features[
+                ~removed_nodes_of_this_type
+            ]
+            original_ids_dict[node_type] = np.where(
+                ~removed_nodes_of_this_type
+            )[0]
+
+        return removed_nodes_dict, original_ids_dict
+
+    @property
+    def unwrapped(self) -> EnvType:
+        """Returns the unwrapped environment."""
+        return self.env  # type: ignore[return-value]

job_shop_lib/reinforcement_learning/_single_job_shop_graph_env.py

@@ -243,8 +243,27 @@ class SingleJobShopGraphEnv(gym.Env):
         *,
         seed: Optional[int] = None,
         options: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[ObservationDict, dict]:
-        """Resets the environment."""
+    ) -> Tuple[ObservationDict, dict[str, Any]]:
+        """Resets the environment.
+
+        Args:
+            seed:
+                Added to match the signature of the parent class. It is not
+                used in this method.
+            options:
+                Additional options to pass to the environment. Not used in
+                this method.
+
+        Returns:
+            A tuple containing the following elements:
+
+            - The observation of the environment.
+            - A dictionary with additional information, keys
+                include: "feature_names", the names of the features in the
+                observation; and "available_operations_with_ids", a list of
+                available a list of available actions in the form of
+                (operation_id, machine_id, job_id).
+        """
         super().reset(seed=seed, options=options)
         self.dispatcher.reset()
         obs = self.get_observation()

job_shop_lib/reinforcement_learning/_utils.py

@@ -1,6 +1,6 @@
 """Utility functions for reinforcement learning."""
 
-from typing import TypeVar, Any, Type, Literal
+from typing import TypeVar, Any, Type
 
 import numpy as np
 from numpy.typing import NDArray
@@ -91,8 +91,10 @@ def add_padding(
 
 
 def create_edge_type_dict(
-    edge_index: NDArray[T], type_ranges: dict[str, tuple[int, int]]
-) -> dict[tuple[str, Literal["to"], str], NDArray[T]]:
+    edge_index: NDArray[T],
+    type_ranges: dict[str, tuple[int, int]],
+    relationship: str = "to",
+) -> dict[tuple[str, str, str], NDArray[T]]:
     """Organizes edges based on node types.
 
     Args:
@@ -101,17 +103,19 @@ def create_edge_type_dict(
         type_ranges: dict[str, tuple[int, int]]
             Dictionary mapping type names to their corresponding index ranges
             [start, end) in the ``edge_index`` array.
+        relationship:
+            A string representing the relationship type between nodes.
 
     Returns:
-        A dictionary with keys (type_i, "to", type_j) and values as edge
-        indices
+        A dictionary with keys (type_i, relationship, type_j) and values as
+        edge indices
     """
-    edge_index_dict: dict[tuple[str, Literal["to"], str], NDArray] = {}
+    edge_index_dict: dict[tuple[str, str, str], NDArray] = {}
     for type_name_i, (start_i, end_i) in type_ranges.items():
         for type_name_j, (start_j, end_j) in type_ranges.items():
-            key: tuple[str, Literal["to"], str] = (
+            key: tuple[str, str, str] = (
                 type_name_i,
-                "to",
+                relationship,
                 type_name_j,
             )
             # Find edges where source is in type_i and target is in type_j
@@ -126,6 +130,40 @@ def create_edge_type_dict(
     return edge_index_dict
 
 
+def map_values(array: NDArray[T], mapping: dict[int, int]) -> NDArray[T]:
+    """Maps values in an array using a mapping.
+
+    Args:
+        array:
+            An NumPy array.
+
+    Returns:
+        A NumPy array where each element has been replaced by its
+        corresponding value from the mapping.
+
+    Raises:
+        ValidationError:
+            If the array contains values that are not in the mapping.
+
+    Examples:
+        >>> map_values(np.array([1, 2, 3]), {1: 10, 2: 20, 3: 30})
+        array([10, 20, 30])
+
+        >>> map_values(np.array([1, 2]), {1: 10, 2: 10, 3: 30})
+        array([10, 10])
+
+    """
+    if array.size == 0:
+        return array
+    try:
+        vectorized_mapping = np.vectorize(mapping.get)
+        return vectorized_mapping(array)
+    except TypeError as e:
+        raise ValidationError(
+            "The array contains values that are not in the mapping."
+        ) from e
+
+
 if __name__ == "__main__":
     import doctest
 

{job_shop_lib-1.0.0b3.dist-info → job_shop_lib-1.0.0b5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: job-shop-lib
-Version: 1.0.0b3
+Version: 1.0.0b5
 Summary: An easy-to-use and modular Python library for the Job Shop Scheduling Problem (JSSP)
 License: MIT
 Author: Pabloo22
@@ -60,7 +60,7 @@ See [this](https://colab.research.google.com/drive/1XV_Rvq1F2ns6DFG8uNj66q_rcoww
 Version 1.0.0 is currently in beta stage and can be installed with:
 
 ```bash
-pip install job-shop-lib==1.0.0b3
+pip install job-shop-lib==1.0.0b5
 ```
 
 Although this version is not stable and may contain breaking changes in subsequent releases, it is recommended to install it to access the new reinforcement learning environments and familiarize yourself with new changes (see the [latest pull requests](https://github.com/Pabloo22/job_shop_lib/pulls?q=is%3Apr+is%3Aclosed)). There is a [documentation page](https://job-shop-lib.readthedocs.io/en/latest/) for versions 1.0.0a3 and onward.

{job_shop_lib-1.0.0b3.dist-info → job_shop_lib-1.0.0b5.dist-info}/RECORD

@@ -1,4 +1,4 @@
-job_shop_lib/__init__.py,sha256=fj4EeaNf1NTAbpB0EH9OnhS_MjP1MCe03Xr6JNrLC_g,643
+job_shop_lib/__init__.py,sha256=3Ip52vXCTCs_JAhhNCpSs1Q2omWPbPzKzM1g47mO5eA,643
 job_shop_lib/_base_solver.py,sha256=p17XmtufNc9Y481cqZUT45pEkUmmW1HWG53dfhIBJH8,1363
 job_shop_lib/_job_shop_instance.py,sha256=hNQGSJj0rEQpS-YhzwWmM6QzCWp6r--89jkghSgLvUs,18380
 job_shop_lib/_operation.py,sha256=hx2atpP8LPj9fvxpZIfhBFr9Uq6JP-MKAX5JzTvFXso,3847
@@ -50,12 +50,13 @@ job_shop_lib/graphs/graph_updaters/__init__.py,sha256=UhnZL55e3cAv7hVetB6bRmIOn8
 job_shop_lib/graphs/graph_updaters/_graph_updater.py,sha256=j1f7iWsa62GVszK2BPaMxnKBCEGWa9owm8g4VWUje8w,1967
 job_shop_lib/graphs/graph_updaters/_residual_graph_updater.py,sha256=SfgmDyMwfW56OBjJPaU76c42IsX5qx9j-eMtrv0DjKk,6047
 job_shop_lib/graphs/graph_updaters/_utils.py,sha256=X5YfwJA1CCgpm1r9C036Gal2CkDh2SSak7wl7TbdjHw,704
-job_shop_lib/reinforcement_learning/__init__.py,sha256=IohAO2-eAdgeAsNG2czRhI5Eu9jKPnDnB8Z-ne_L1as,1124
-job_shop_lib/reinforcement_learning/_multi_job_shop_graph_env.py,sha256=ib1Y6cItVvId4PfesiQ0XKbh9y6h8LVhD0gYDO4wSlk,15732
+job_shop_lib/reinforcement_learning/__init__.py,sha256=opqJyVJ6VPyeaQOQr4hmUTkiUAXOi5tbyCnuNw5jLTI,1421
+job_shop_lib/reinforcement_learning/_multi_job_shop_graph_env.py,sha256=memQefVWqatRNodt8hBXVvFcgQKJRmnuB8AWqiDl8_k,15746
+job_shop_lib/reinforcement_learning/_resource_task_graph_observation.py,sha256=ROu_oTsjLvuU5cdvZ2A8D5HWlEZ3NfPtjNWJYH5ilJA,9648
 job_shop_lib/reinforcement_learning/_reward_observers.py,sha256=iWHccnujeAKyTQn2ilQ4BhcEccoSTyJqQ5yOiP5GG_Y,2984
-job_shop_lib/reinforcement_learning/_single_job_shop_graph_env.py,sha256=DZnXeXmzMGKq-vwFhxukmSDN1UyrkUfbnjpjFtC9_Bs,15845
+job_shop_lib/reinforcement_learning/_single_job_shop_graph_env.py,sha256=3mljeI4k9haLDuDZZhv4NcLpb1_xOlQNdEqMoRsh6bw,16592
 job_shop_lib/reinforcement_learning/_types_and_constants.py,sha256=xozdM_Wabdbe9e1a769p5980OSNBwQqc9yyaSGW2ODQ,1743
-job_shop_lib/reinforcement_learning/_utils.py,sha256=PZWUZVc_Du90PYvKeTxe0dLkCuK3KwQ3D5I54kk5x0U,3808
+job_shop_lib/reinforcement_learning/_utils.py,sha256=giikAj9Xl2f_cYq_AKCrgHn79TRLfGhYNqXz7bEnKuk,4827
 job_shop_lib/visualization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 job_shop_lib/visualization/gantt/__init__.py,sha256=HGXwRgDuMAldqU0JBdiZCd5e79XBz1r96qHeDVlzE54,1145
 job_shop_lib/visualization/gantt/_gantt_chart_creator.py,sha256=LTsVhpB1Fb_2o08HRZPPXSekwzR7fyTSC6h549XMqhU,8638
@@ -64,7 +65,7 @@ job_shop_lib/visualization/gantt/_plot_gantt_chart.py,sha256=9-NSSNsVcW8gYLZtAuF
 job_shop_lib/visualization/graphs/__init__.py,sha256=282hZFg07EyQu4HVt4GzFfYnY6ZF376IMjnWZ5eg0ZQ,611
 job_shop_lib/visualization/graphs/_plot_disjunctive_graph.py,sha256=4VBMYiFXXkCGSnGYN9iqNtWrbLJQxAMHojPHhAbdA0s,14387
 job_shop_lib/visualization/graphs/_plot_resource_task_graph.py,sha256=RgJqHS5hJh3KkyaLbtpG_bER981BFRwGpflz7I7gS64,13271
-job_shop_lib-1.0.0b3.dist-info/LICENSE,sha256=9mggivMGd5taAu3xbmBway-VQZMBzurBGHofFopvUsQ,1069
-job_shop_lib-1.0.0b3.dist-info/METADATA,sha256=m-NzkaDXbblEXE7O-iQEYOdQQCPDvj6LiRcJKqenpHg,16424
-job_shop_lib-1.0.0b3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-job_shop_lib-1.0.0b3.dist-info/RECORD,,
+job_shop_lib-1.0.0b5.dist-info/LICENSE,sha256=9mggivMGd5taAu3xbmBway-VQZMBzurBGHofFopvUsQ,1069
+job_shop_lib-1.0.0b5.dist-info/METADATA,sha256=NxGDA_C8w6GQlV1zcrkIawyHekRWR8VTvNwQR9s7Ci4,16424
+job_shop_lib-1.0.0b5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+job_shop_lib-1.0.0b5.dist-info/RECORD,,