chemrecon 0.1.1__py3-none-any.whl

chemrecon/database/params.py

@@ -0,0 +1,88 @@
+""" This script loads and exports connection params.
+"""
+import os
+from typing import Optional
+
+
+class Params:
+    """
+    Represents database connection parameters and provides a method to generate a formatted
+    connection string compatible with PostgreSQL.
+    """
+    connection_title: str  # Name of the connection file
+    db_name: str
+    db_host: str
+    db_port: str
+    username: str
+    password: str
+
+    def __init__(
+        self,
+        connection_title: str,
+        db_name: str,
+        db_host: str,
+        db_port: str,
+        username: str,
+        password: str
+    ):
+        self.connection_title = connection_title
+        self.db_name = db_name
+        self.db_host = db_host
+        self.db_port = db_port
+        self.username = username
+        self.password = password
+
+    def connection_string(self) -> str:
+        """ Returns a connection string compatible with PostgreSQL."""
+        return f'postgres://{self.username}:{self.password}@{self.db_host}:{self.db_port}/{self.db_name}'
+
+
+# Load params
+parameter_sets: list[Params] = list()
+
+script_path = os.path.dirname(__file__)
+rel_path = 'connection_params/'
+params_dir = os.path.join(script_path, rel_path)
+for param_file in os.listdir(params_dir):
+    try:
+        with open(os.path.join(params_dir, param_file)) as f:
+            flines = f.readlines()
+            param = Params(
+                param_file,
+                db_name = flines[0].strip(),
+                db_host = flines[3].strip(),
+                db_port = flines[4].strip(),
+                username = flines[1].strip(),
+                password = flines[2].strip()
+            )
+            if len(flines) > 5:
+                raise ValueError('Malformed database connection parameters file.')
+
+            parameter_sets.append(param)
+    except (KeyError, IndexError):
+        print('Invalid connection parameters file.')
+        pass
+
+# Set params to find easily
+
+# Test
+local_docker_init: Optional[Params] = None
+local_docker_dev: Optional[Params] = None
+local_docker_pub: Optional[Params] = None
+
+# Production
+chemrecon_dev: Optional[Params] = None
+chemrecon_pub: Optional[Params] = None
+
+for p in parameter_sets:
+    match p.connection_title:
+        case 'local_docker_init.dbinfo':
+            local_docker_init = p
+        case 'local_docker_dev.dbinfo':
+            local_docker_dev = p
+        case 'local_docker_pub.dbinfo':
+            local_docker_pub = p
+        case 'chemrecon_dev.dbinfo':
+            chemrecon_dev = p
+        case 'chemrecon_pub.dbinfo':
+            chemrecon_pub = p

chemrecon/entrygraph/draw.py

@@ -0,0 +1,119 @@
+""" Drawing utilites for each type of entry.
+"""
+from __future__ import annotations
+
+from typing import Type, Optional
+
+from PIL import Image
+import rustworkx.visualization as rx_vis
+
+import chemrecon.entrygraph.entrygraph
+from chemrecon.schema import Entry, Relation
+
+attr_none: str = "''"
+linewidth: int = 30
+
+
+# noinspection PyProtectedMember
+class EntryGraphDrawer():
+    entrygraph_type: Type[chemrecon.entrygraph.entrygraph.EntryGraph]
+
+    # Draw settings
+    draw_reconids: bool
+
+    def __init__(
+        self,
+        # Settings
+        draw_reconids: bool = False
+    ):
+        """ The initializer for the Drawer sets the settings
+        """
+        self.draw_reconids = draw_reconids
+
+    def draw(
+        self,
+        entrygraph: chemrecon.entrygraph.entrygraph.EntryGraph,
+        filename: str = None,
+        filetype: str = 'png',
+        scores: Optional[dict[Entry, float]] = None
+    ) -> Image.Image:
+        """ Draw the entry graph.
+
+            For possible file formats, see the the corresponding
+            [RustWorkX documentation](https://www.rustworkx.org/apiref/rustworkx.visualization.graphviz_draw.html).
+        """
+
+        g_ = entrygraph.g.copy()
+
+        # Empty score dict if not given
+        score_dict: dict[Entry, float] = scores if scores is not None else dict()
+
+        # Remove duplicate edges (for symmetric relations)
+        for e_index in g_.edge_indices():
+            data = g_.get_edge_data_by_index(e_index)
+            if data.relation.symmetric:
+                x, y = g_.get_edge_endpoints_by_index(e_index)
+                if x > y:
+                    g_.remove_edge_from_index(e_index)
+
+        # Remove self-edges
+        for e_index in g_.edge_indices():
+            x, y = g_.get_edge_endpoints_by_index(e_index)
+            if x == y:
+                g_.remove_edge_from_index(e_index)
+
+        return rx_vis.graphviz_draw(
+            g_,
+            node_attr_fn = lambda v: self.node_attr(v.entry, score_dict.get(v.entry, None)),
+            edge_attr_fn = lambda e: self.edge_attr(e.relation),
+            graph_attr = None,  # TODO
+            filename = filename,
+            image_type = filetype,
+        )
+
+    def node_attr(self, entry: Entry, score: Optional[float]) -> dict[str, str]:
+        assert entry.recon_id is not None
+
+        # Label
+        if self.draw_reconids:
+            recon_id_part = f'\n{entry.recon_id}'
+        else:
+            recon_id_part = ''
+
+        # Break lines of label
+        entrylabel = entry._vis_str()
+        label_lines = entrylabel.splitlines()
+        split_label: list[str] = list()
+        for l in label_lines:
+            while l != '':
+                split_label.append(l[:linewidth])
+                l = l[linewidth:]
+        entry_label_final = ('\n').join(split_label)
+
+        # Determine label
+        label: str = f'{entry.get_table_name()}\n{entry_label_final}{recon_id_part}'
+        if score is not None:
+            label += f'\n{score:.3f}'
+
+        # Other drawing
+        d = {
+            'style': 'filled',
+            'shape': 'box',
+            'label': label
+        }
+        d.update(entry._vis_attrs())
+        return d
+
+    def edge_attr(self, rel: Relation) -> dict[str, str]:
+        label = rel._vis_str()
+        d = {
+            'label': label,
+            'dir': 'none' if rel.symmetric else 'forward',
+        }
+
+        # Update with rel_type-specific attributes and return
+        d.update(rel._vis_attrs())
+        return d
+
+# Default drawer - this is used by the EntryGraph .draw() and .show() methods.
+defaultdrawer = EntryGraphDrawer()

chemrecon/entrygraph/entrygraph.py

@@ -0,0 +1,301 @@
+from __future__ import annotations
+
+from typing import Optional, ClassVar, Generator
+
+import rustworkx as rx
+
+from chemrecon.entrygraph.draw import defaultdrawer
+from chemrecon.entrygraph.explore_procedure import ExploreProcedure
+from chemrecon.schema import Entry, Relation
+
+from chemrecon import connection as connection
+
+type VertexIndex = int
+type EdgeIndex = int
+type ReconID = int
+
+
+# Vertices and Edges
+# --------------------------------------------------------------------------------------------------------------
+# Vertex and Edge objects are used as the payload for the RustWorkX graph.
+
+class VertexBase:
+    pass
+
+
+class Vertex[E: Entry](VertexBase):
+    """ A database entry in a graph context.
+        None is used as the type parameter for artificial vertices (those not corresponding to a DB object).
+        A vertex is initial if its generation is 0.
+    """
+    entry: E  #: The entry represented by this vertex
+    recon_id: ReconID  #:
+    vertex_index: VertexIndex  #: Index in the graph representation.
+    generation: int  #: Distance from initial entries.
+
+    def __init__(self, entry: E, generation: int):
+        assert entry.recon_id is not None
+        self.entry = entry
+        self.recon_id = entry.recon_id
+        self.generation = generation
+
+
+class EdgeBase:
+    pass
+
+
+class Edge[R: Relation | None](EdgeBase):
+    """ A database relation in a graph context.
+        None is used as the type parameter for artificial edges (those not corresponding to a DB object).
+    """
+    relation: R  #:
+    recon_id_1: ReconID  #:
+    recon_id_2: ReconID  #:
+    edge_index: EdgeIndex  #: Index in the graph representation
+
+    def __init__(self, relation: R):
+        self.relation = relation
+        self.recon_id_1 = relation.recon_id_1
+        self.recon_id_2 = relation.recon_id_2
+
+
+# Special case graph vertices and edges which do not represent database entries or relations
+class SourceVertexArtificial(VertexBase):
+    """ Artificial vertex used as a source"""
+
+    def __init__(self):
+        pass
+
+
+class SourceEdgeArtificial(EdgeBase):
+    """ Artificial edge from the SourceVertex to initial vertices in the graph. """
+
+    def __init__(self):
+        pass
+
+
+class ReturnEdgeArtificial(EdgeBase):
+    """ Temporary edge added from sink vertices back to the SourceVertex for pagerank purposes. """
+
+    def __init__(self):
+        pass
+
+
+# Generic EntryGraph
+# --------------------------------------------------------------------------------------------------------------
+class EntryGraph:
+    """
+    Represents a directed graph structure of entries and relations.
+
+    :ivar g: RustWorkX directed graph object used to store vertices and edges.
+    :type g: rx.PyDiGraph
+    :ivar index_entry: Maps Entry objects to corresponding vertex indices.
+    :type index_entry: dict[Entry, VertexIndex]
+    :ivar index_reconid: Maps (Entry type, ReconID) tuples to vertex indices.
+    :type index_reconid: dict[tuple[type[Entry], ReconID], VertexIndex]
+    :ivar index_relation: Maps Relation objects to corresponding edge indices.
+    :type index_relation: dict[Relation, EdgeIndex]
+    :ivar initial_entries: Set of Entry objects that serve as the initial vertices in the graph.
+    :type initial_entries: set[Entry]
+    :ivar initial_vertices: Set of vertex indices corresponding to the initial entries.
+    :type initial_vertices: set[VertexIndex]
+    :ivar parent_entrygraph: Reference to the parent EntryGraph if this graph
+        is a subgraph view.
+    :type parent_entrygraph: Optional[EntryGraph]
+    """
+    g: rx.PyDiGraph
+
+    # Indices for looking up vertices / edges based on the underlying entries
+    index_entry: dict[Entry, VertexIndex]
+    index_reconid: dict[tuple[type[Entry], ReconID], VertexIndex]
+    index_relation: dict[Relation, EdgeIndex]
+
+    # Misc
+    initial_entries: set[Entry]
+    initial_vertices: set[VertexIndex]
+
+    # Parentage
+    parent_entrygraph: Optional[EntryGraph]  # If created as a subgraph view, gives the parent graph
+
+    # List of finished ExplorationProcedure objects to call when exploring
+    explore_procedures: ClassVar[list[ExploreProcedure]]
+    post_explore_procedures: ClassVar[list[ExploreProcedure]]
+    transitive_subprocedures: ClassVar[dict[type[Relation], tuple[ExploreProcedure, ExploreProcedure]]]
+
+    # TODO scoring as abstract class variables
+
+
+    def __init__(
+        self,  # TODO specification refactor
+        initial_entries: set[Entry]
+    ):
+        """ Initialize an entrygraph with a set of initial vertices.
+        """
+        # Consistency check
+        # TODO init vertices in allowed types
+
+        # Create
+        self.index_entry = dict()
+        self.index_reconid = dict()
+        self.index_relation = dict()
+
+        # Create graph
+        self.g = rx.PyDiGraph(multigraph = True)
+
+        # Add initial vertices
+        # The initial vertices may not have recon-ids attached, so we add those by DB lookup
+        for entry in initial_entries:
+            if entry.recon_id is None:
+                res = connection.handler.get_entry_by_index(entry)
+                if res is None:
+                    continue
+                entry.recon_id = res.recon_id
+
+        # If no initial vertices have recon ids, abort
+        if all(e.recon_id is None for e in initial_entries):
+            raise ValueError('No initial entries have recon ids.')
+
+        self.initial_entries = initial_entries
+        self.initial_vertices = {
+            self.add_vertex(e, generation = 0)
+            for e in initial_entries
+        }
+
+    # Getters
+    # ----------------------------------------------------------------------------------------------------------
+    def get_vertex_index_by_recon_id(self, entry_type: type[Entry], recon_id: ReconID) -> Optional[VertexIndex]:
+        """
+        Retrieve the vertex index associated with the given entry type and recon ID.
+        If no matching vertex index is found, the method returns None.
+
+        :param entry_type: The class type of the entry being queried.
+        :type entry_type: type[Entry]
+        :param recon_id: The recon ID that uniquely identifies the entry.
+        :type recon_id: ReconID
+        :return: The vertex index corresponding to the given entry type and recon ID, or None
+                 if not found.
+        :rtype: Optional[VertexIndex]
+        """
+        return self.index_reconid.get((entry_type, recon_id), None)
+
+    def get_vertex_by_entry(self, entry: Entry) -> Optional[VertexIndex]:
+        """
+        Retrieves the vertex index associated with a given entry and returns None if it does not exist.
+
+        :param entry: The entry object to search for in the index map.
+        :type entry: Entry
+        :return: The vertex index associated with the given entry, or None if the entry
+                 does not exist in the mapping.
+        :rtype: Optional[VertexIndex]
+        """
+        return self.index_entry.get(entry, None)
+
+    def get_vertex_by_vertex_index(self, vertex_index: VertexIndex) -> Optional[Vertex]:
+        """
+        Retrieve the vertex corresponding to the given vertex index, returns None if it does not exist.
+        :param vertex_index: Index of the vertex to retrieve.
+        :type vertex_index: VertexIndex
+        :return: The vertex associated with the given index or None if it does not exist.
+        :rtype: Optional[Vertex]
+        """
+        return self.g.get_node_data(vertex_index)
+
+    def get_out_edges_of_vertex(self, vertex_index: VertexIndex) -> list[tuple[Edge, Vertex]]:
+        """
+        Retrieve all outgoing edges of a vertex along with the corresponding target
+        vertices.
+
+        This method provides a list of tuples where each tuple contains the edge data
+        and the target vertex data for the edges originating from the specified
+        vertex. This function is useful for analysing the outgoing connections
+        and relationships in a graph data structure.
+
+        :param vertex_index: Index of the vertex whose outgoing edges are to be queried.
+        :type vertex_index: VertexIndex
+        :return: A list of tuples with each tuple containing edge data and the target
+                 vertex data.
+        :rtype: list[tuple[Edge, Vertex]]
+        """
+        return [
+            (edge_data, self.g.get_node_data(target_index))
+            for _, target_index, edge_data in self.g.out_edges(vertex_index)
+        ]
+
+    # Iterators
+    def vertices(self) -> Generator[Vertex, None, None]:
+        """ Generator over vertices (entries). """
+        yield from self.g.nodes()
+
+    def edges(self) -> Generator[Edge, None, None]:
+        """ Generator over edges (relations). """
+        yield from self.g.edges()
+
+    # Adders
+    # ----------------------------------------------------------------------------------------------------------
+    def add_vertex(self, entry: Entry, generation: int = -1) -> VertexIndex:
+        """ Add a new entry to the graph. If it already exists, return the index.
+        """
+        if entry.recon_id is None:
+            raise ValueError('Cannot construct vertex from entry without reconid.')
+        lookup = self.get_vertex_index_by_recon_id(type(entry), entry.recon_id)
+        if lookup is not None:
+            return lookup
+
+        # Add the vertex
+        vertex = Vertex(entry, generation = generation)
+        v_index: VertexIndex = self.g.add_node(vertex)
+        self.index_reconid[(type(entry), entry.recon_id)] = v_index
+        self.index_entry[entry] = v_index
+        vertex.vertex_index = v_index
+        return v_index
+
+    def add_edge(
+        self,
+        source_v_index: VertexIndex,
+        target_v_index: VertexIndex,
+        relation: Relation
+    ) -> EdgeIndex:
+        """ Add a relation between the specified source and target vertices
+        """
+        # Try returning existing edge
+        try:
+            return self.index_relation[relation]
+        except KeyError:
+            # Does not already exist
+            edge = Edge(relation)
+            edge_index: EdgeIndex = self.g.add_edge(source_v_index, target_v_index, edge)
+            edge.edge_index = edge_index
+            self.index_relation[relation] = edge_index
+            return edge_index
+
+    def add_vertex_from[T1: Entry, T2: Entry](
+        self,
+        from_index: VertexIndex,
+        relation: Relation[T1, T2],
+        entry: T2,
+        generation: int
+    ) -> VertexIndex:
+        """ Adds a new entry vertex along with a relation. Returns vertex_index of the newly creted vertex.
+            If the vertex already exists, return add the edge and return the vertex index.
+        """
+        target_v_index = self.add_vertex(entry, generation = generation)
+        self.add_edge(from_index, target_v_index, relation)
+        return target_v_index
+
+    # Draw
+    # ----------------------------------------------------------------------------------------------------------
+    def draw(
+        self,
+        filename: str = None,
+        filetype: str = 'jpg',
+        scores: Optional[dict[Entry, float]] = None
+    ):
+        """ Draw using default settings, returning an image, or write to the disk if the filename is specified.
+            For more settings, create an EntryGraphDrawer instance.
+        """
+        return defaultdrawer.draw(self, filename = filename, filetype = filetype, scores = scores)
+
+    def show(self, scores: Optional[dict[Entry, float]] = None):
+        """ Same as draw, but displays the image in a window instead of saving to disk.
+        """
+        return self.draw(scores = scores).show()