Merge pull request #36 from JaxGaussianProcesses/restructure

Restructure modules

.pre-commit-config.yaml

@@ -18,4 +18,18 @@ repos:
         language: system
         types: [python]
         require_serial: true
-
+      - id: autoflake
+        name: autoflake
+        entry: autoflake
+        language: python
+        "types": [python]
+        require_serial: true
+        args:
+          - "--in-place"
+          - "--expand-star-imports"
+          - "--remove-all-unused-imports"
+          - "--remove-unused-variables"
+  - repo: https://github.com/psf/black
+    rev: 22.3.0
+    hooks:
+      - id: black

README.md

@@ -3,3 +3,7 @@
 
 [![codecov](https://codecov.io/gh/JaxGaussianProcesses/JaxKern/branch/main/graph/badge.svg?token=8WD7YYMPFS)](https://codecov.io/gh/JaxGaussianProcesses/JaxKern)
 [![CircleCI](https://dl.circleci.com/status-badge/img/gh/JaxGaussianProcesses/JaxKern/tree/main.svg?style=svg)](https://dl.circleci.com/status-badge/redirect/gh/JaxGaussianProcesses/JaxKern/tree/main)
+[![Documentation Status](https://readthedocs.org/projects/gpjax/badge/?version=latest)](https://gpjax.readthedocs.io/en/latest/?badge=latest)
+[![PyPI version](https://badge.fury.io/py/jaxkern.svg)](https://badge.fury.io/py/jaxkern)
+[![Downloads](https://pepy.tech/badge/jaxkern)](https://pepy.tech/project/jaxkern)
+[![Slack Invite](https://img.shields.io/badge/Slack_Invite--blue?style=social&logo=slack)](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw)

jaxkern/__init__.py

@@ -1,30 +1,57 @@
+# Copyright 2022 The JaxGaussianProcesses Contributors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
 """JaxKern."""
-from .kernels import (
+from .base import ProductKernel, SumKernel
+from .computations import (
+    ConstantDiagonalKernelComputation,
+    DenseKernelComputation,
+    DiagonalKernelComputation,
+    EigenKernelComputation,
+)
+from .nonstationary import Linear, Polynomial
+from .stationary import (
     RBF,
-    GraphKernel,
     Matern12,
     Matern32,
     Matern52,
-    Polynomial,
-    ProductKernel,
-    SumKernel,
-    DenseKernelComputation,
-    DiagonalKernelComputation,
-    ConstantDiagonalKernelComputation,
+    RationalQuadratic,
+    Periodic,
+    PoweredExponential,
+    White,
 )
+from .non_euclidean import GraphKernel
 
 __all__ = [
     "RBF",
     "GraphKernel",
     "Matern12",
     "Matern32",
     "Matern52",
+    "Linear",
     "Polynomial",
     "ProductKernel",
     "SumKernel",
     "DenseKernelComputation",
     "DiagonalKernelComputation",
     "ConstantDiagonalKernelComputation",
+    "EigenKernelComputation",
+    "PoweredExponential",
+    "Periodic",
+    "RationalQuadratic",
+    "White",
 ]
 
 from . import _version