Merge pull request 'change full context to o3-mini' (#14) from impl_test into main

Reviewed-on: #14

.gitea/scripts/code_review.py

@@ -0,0 +1,304 @@
+"""Code Reviewer for Gitea."""
+
+import asyncio
+import fnmatch
+import json
+import os
+import re
+from typing import Any
+
+import requests
+import aiohttp
+from model import Model
+
+ACCESS_TOKEN = os.getenv("ACCESS_TOKEN", "")
+HEADERS = {"Authorization": f"token {ACCESS_TOKEN}"}
+
+GITHUB_EVENT_PATH = os.getenv("GITHUB_EVENT_PATH")
+try:
+    with open(GITHUB_EVENT_PATH, "r") as f:
+        EVENT_DATA = json.load(f)
+except FileNotFoundError:
+    print("Failed to load event data.")
+    exit(1)
+
+FULL_CONTEXT_MODEL_NAME = os.getenv("FULL_CONTEXT_MODEL", "")
+SINGLE_CHUNK_MODEL_NAME = os.getenv("SINGLE_CHUNK_MODEL", "")
+FULL_CONTEXT_API_KEY = os.getenv("FULL_CONTEXT_API_KEY", "")
+SINGLE_CHUNK_API_KEY = os.getenv("SINGLE_CHUNK_API_KEY", "")
+
+EXCLUDE_PATTERNS = os.getenv("EXCLUDE", "").split(",")
+
+
+def get_diff() -> str | None:
+    """Get code difference between base and head from Gitea.
+
+    Returns:
+        str | None: code difference between base and head, or None if failed to get diff
+    """
+    url = EVENT_DATA["pull_request"]["diff_url"]
+    try:
+        response = requests.get(url, headers=HEADERS)
+        response.raise_for_status()
+        return response.text
+    except requests.RequestException as e:
+        print(f"Failed to get diff: {e}")
+        return None
+
+
+def parse_diff(diff: str) -> list[dict[str, Any]]:
+    """Parse diff into list of dicts.
+
+    Args:
+        diff: str, code difference between base and head
+
+    Returns:
+        list[dict[str, Any]]: list of dicts, each dict represents a code chunks
+    """
+    file_pattern = re.compile(
+        r"(?s)diff --git a/(.+?) b/(.*?)\r?\n(.*?)(?=diff --git a/|$)", re.S
+    )
+    old_new_pattern = re.compile(r"(?m)^(---|\+\+\+)\s+(.*)$")
+    chunk_range_pattern = re.compile(
+        r"@@ -(\d+)(?:,(\d+))? \+(\d+)(?:,(\d+))? @@(.*?)?(?=@@|\Z)",
+        re.MULTILINE | re.DOTALL,
+    )
+    list_diff = []
+    for match in file_pattern.finditer(diff):
+        diff_text = match.group(3)
+
+        old_new_match = list(old_new_pattern.finditer(diff_text))
+        if len(old_new_match) != 2:
+            continue
+
+        old_file = old_new_match[0].group(2)
+        old_file = old_file.lstrip("a/") if old_file.startswith("a/") else old_file
+
+        new_file = old_new_match[1].group(2)
+        if new_file == "/dev/null":
+            print("Neglict deleted file")
+            continue
+        new_file = new_file.lstrip("b/")
+        if any(fnmatch.fnmatch(new_file, pattern) for pattern in EXCLUDE_PATTERNS):
+            print(f"Exclude file {new_file}")
+            continue
+
+        output_diff_text = []
+        for chunk_range_match in chunk_range_pattern.finditer(diff_text):
+            old_idx = int(chunk_range_match.group(1))
+            new_idx = int(chunk_range_match.group(3))
+            for line in chunk_range_match.group(5).splitlines():
+                if line.startswith("-"):
+                    output_diff_text.append(f"{old_idx} None {line}")
+                    old_idx += 1
+                elif line.startswith("+"):
+                    output_diff_text.append(f"None {new_idx} {line}")
+                    new_idx += 1
+                else:
+                    output_diff_text.append(f"{old_idx} {new_idx} {line}")
+                    old_idx += 1
+                    new_idx += 1
+
+        output_diff_text = "\n".join(output_diff_text)
+        list_diff.append(
+            {
+                "file": new_file,
+                "chunk": output_diff_text,
+            }
+        )
+    return list_diff
+
+
+def create_comment(
+    file: str, ai_response: list[dict[str, Any]]
+) -> list[dict[str, Any]]:
+    """Create comments for single chunk review.
+
+    Args:
+        file: str, file name
+        ai_response: list[dict[str, Any]], AI response for single chunk review
+
+    Returns:
+        list[dict[str, Any]]: comments for single chunk review
+    """
+    comments = []
+    for ai_response in ai_response:
+        comments.append(
+            {
+                "body": f"[REVIEW] {ai_response['reviewComment']}",
+                "path": file,
+                "new_position": int(ai_response["lineNumber"]),
+            }
+        )
+    return comments
+
+
+async def analyze_single_chunks(
+    single_chunk_model: Model, parsed_diff: list[dict[str, Any]]
+) -> list[dict[str, Any]]:
+    """Analyze single chunks and create comments.
+
+    Args:
+        single_chunk_model: AI Session for single chunk analysis
+        parsed_diff: list[dict[str, Any]], parsed diff
+
+    Returns:
+        list[dict[str, Any]]: comments for single chunk review
+    """
+
+    async def process_single_chunk(diff: dict[str, Any]):
+        file = diff["file"]
+        chunk = diff["chunk"]
+        response = await single_chunk_model.get_response_single_chunk(
+            file, title, description, chunk
+        )
+        response = response.strip("`").lstrip("json").strip() or "[]"
+
+        try:
+            response_json = json.loads(response)
+            return create_comment(file, response_json)
+        except json.JSONDecodeError:
+            print(f"Failed to parse response: {response}")
+            return []
+
+    title = EVENT_DATA["pull_request"]["title"]
+    description = EVENT_DATA["pull_request"]["body"]
+    tasks = [process_single_chunk(diff) for diff in parsed_diff]
+    results = await asyncio.gather(*tasks)
+
+    # Flatten the list of comments
+    comments = [comment for result in results for comment in result]
+    return comments
+
+
+async def get_file_content(file: str) -> str | None:
+    """Get file content from Gitea.
+
+    Args:
+        file: str, file name
+
+    Returns:
+        str | None: file content, or None if failed to get file content
+    """
+    repo_url = EVENT_DATA["pull_request"]["head"]["repo"]["url"]
+    branch = EVENT_DATA["pull_request"]["head"]["ref"]
+
+    replaced_file = file.replace("/", "%2F")
+    url = f"{repo_url}/raw/{branch}%2F{replaced_file}?ref={branch}"
+
+    try:
+        async with aiohttp.ClientSession(headers=HEADERS) as session:
+            async with session.get(url) as response:
+                response.raise_for_status()
+                return await response.text()
+    except aiohttp.ClientError as e:  # More specific exception handling
+        print(f"Network error fetching {file}: {e}")
+    except asyncio.TimeoutError:
+        print(f"Timeout fetching {file}")
+    return None
+
+
+async def analyze_full_context(
+    full_context_model: Model, parsed_diff: list[dict[str, Any]]
+) -> str:
+    """Analyze full context and create review.
+
+    Args:
+        full_context_model: AI Session for full context analysis
+        parsed_diff: list[dict[str, Any]], parsed diff
+
+    Returns:
+        str: review for full context
+    """
+
+    async def get_file_data(diff: dict[str, Any]):
+        file = diff["file"]
+        chunk = diff["chunk"]
+        content = get_file_content(file)
+        if content is None:
+            return None
+        return f"File: {file}\n{content}\nDiff: {chunk}"
+
+    tasks = [get_file_data(diff) for diff in parsed_diff]
+    file_contents_list = await asyncio.gather(*tasks)
+
+    file_contents = [item for item in file_contents_list if item is not None]
+
+    if not file_contents:
+        return ""
+
+    title = EVENT_DATA["pull_request"]["title"]
+    description = EVENT_DATA["pull_request"]["body"]
+    response = await full_context_model.get_response_full_context(
+        title, description, file_contents
+    )
+    response = response.strip("`").lstrip("markdown").strip()
+    return response
+
+
+def post_review(
+    full_context_review: str, single_chunk_comments: list[dict[str, Any]]
+) -> None:
+    """Post review to Gitea.
+
+    Args:
+        full_context_review: str, review for full context
+        single_chunk_comments: list[dict[str, Any]], comments for single chunk review
+    """
+    repo_url = EVENT_DATA["pull_request"]["head"]["repo"]["url"]
+    pull_number = EVENT_DATA["number"]
+    commit_id = EVENT_DATA["pull_request"]["head"]["sha"]
+    url = f"{repo_url}/pulls/{pull_number}/reviews"
+    data = {
+        "body": full_context_review,
+        "event": "COMMENT",
+        "comments": single_chunk_comments,
+        "commit_id": commit_id,
+    }
+    response = requests.post(url, headers=HEADERS, json=data)
+    response.raise_for_status()
+
+
+async def main() -> None:
+    """Code Reviewer for Gitea: Asynchronous version."""
+    if EVENT_DATA["action"] not in ["opened", "synchronized"]:
+        print("Unsupported event.")
+        return
+
+    diff = get_diff()
+    if diff is None:
+        return
+    elif not diff:
+        print("No diff found.")
+        return
+
+    full_context_model = Model(
+        model=FULL_CONTEXT_MODEL_NAME,
+        api_key=FULL_CONTEXT_API_KEY,
+        is_full_context=True,
+    )
+    single_chunk_model = Model(
+        model=SINGLE_CHUNK_MODEL_NAME,
+        api_key=SINGLE_CHUNK_API_KEY,
+        is_full_context=False,
+    )
+
+    parsed_diff = parse_diff(diff)
+    comments_task = asyncio.create_task(
+        analyze_single_chunks(single_chunk_model, parsed_diff)
+    )
+
+    if EVENT_DATA["action"] == "opened":
+        full_context_response_task = asyncio.create_task(
+            analyze_full_context(full_context_model, parsed_diff)
+        )
+        full_context_response = await full_context_response_task
+    else:
+        full_context_response = ""
+
+    comments = await comments_task
+    post_review(full_context_response, comments)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

.gitea/scripts/model.py

@@ -0,0 +1,261 @@
+"""Model for code review."""
+
+from enum import Enum
+from typing import Any
+
+import google.generativeai as genai
+import typing_extensions as typing
+from anthropic import AsyncAnthropic
+from openai import AsyncOpenAI
+
+
+class GoogleResponse(typing.TypedDict):
+    """The response from Google model."""
+
+    lineNumber: int
+    reviewComment: str
+
+
+class ModelProvider(Enum):
+    """The model provider."""
+
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+    GOOGLE = "google"
+    DEEPSEEK = "deepseek"
+
+    @classmethod
+    def from_model(cls, model: str) -> "ModelProvider":
+        """Get the model provider from the model name.
+
+        Args:
+            model (str): The model name.
+
+        Returns:
+            ModelProvider: The model provider.
+        """
+        for prefix, provider in PREFIX_TO_MODEL.items():
+            if model.startswith(prefix):
+                return provider
+        raise ValueError(f"Unknown model: {model}")
+
+
+PREFIX_TO_MODEL = {
+    "gpt": ModelProvider.OPENAI,
+    "o1": ModelProvider.OPENAI,
+    "o3": ModelProvider.OPENAI,
+    "claude": ModelProvider.ANTHROPIC,
+    "gemini": ModelProvider.GOOGLE,
+    "deepseek": ModelProvider.DEEPSEEK,
+}
+
+
+class Model:
+    """The model class.
+
+    Attributes:
+        model (str): The model name.
+        api_key (str): The API key.
+        system_prompt (str): The system prompt.
+        max_tokens (int): The maximum tokens.
+    """
+
+    def __init__(  # noqa: D107
+        self,
+        model: str,
+        api_key: str,
+        is_full_context: bool,
+        max_tokens: int = 4196,
+    ):
+        self.model = model
+        self.system_prompt = (
+            FULL_CONTEXT_SYSTEM_PROMPT
+            if is_full_context
+            else SINGLE_CHUNK_SYSTEM_PROMPT
+        )
+        self.max_tokens = max_tokens
+        self.provider = ModelProvider.from_model(model)
+        self.session = self.create_session(api_key)
+
+    def create_session(self, api_key: str) -> Any:
+        """Create a session for the model.
+
+        Args:
+            api_key (str): The API key.
+
+        Returns:
+            Any: The session.
+        """
+        match self.provider:
+            case ModelProvider.OPENAI:
+                return AsyncOpenAI(api_key=api_key)
+            case ModelProvider.ANTHROPIC:
+                return AsyncAnthropic(api_key=api_key)
+            case ModelProvider.GOOGLE:
+                genai.configure(api_key=api_key)
+                return genai.GenerativeModel(
+                    model_name=self.model, system_instruction=self.system_prompt
+                )
+            case ModelProvider.DEEPSEEK:
+                return AsyncOpenAI(api_key=api_key, base_url="https://api.deepseek.com")
+
+    async def request(self, prompt: str) -> str:
+        """Request the model to generate a response.
+
+        Args:
+            prompt (str): The prompt to generate a response for.
+
+        Returns:
+            str: The generated response.
+        """
+        match self.provider:
+            case ModelProvider.OPENAI | ModelProvider.DEEPSEEK:
+                response = await self.session.chat.completions.create(
+                    model=self.model,
+                    messages=[
+                        {"role": "system", "content": self.system_prompt},
+                        {"role": "user", "content": prompt},
+                    ],
+                    temperature=0.2,
+                    max_tokens=self.max_tokens,
+                    top_p=1,
+                    frequency_penalty=0,
+                    presence_penalty=0,
+                )
+                return response.choices[0].message.content.strip()
+            case ModelProvider.ANTHROPIC:
+                response = await self.session.messages.create(
+                    model=self.model,
+                    messages=[{"role": "user", "content": prompt}],
+                    system=[
+                        {
+                            "type": "text",
+                            "text": self.system_prompt,
+                            "cache_control": {"type": "ephemeral"},
+                        }
+                    ],
+                    temperature=0.2,
+                    max_tokens=self.max_tokens,
+                )
+                return response.content[0].text.strip()
+            case ModelProvider.GOOGLE:
+                response = await self.session.generate_content_async(
+                    prompt,
+                    generation_config=genai.GenerationConfig(
+                        response_mime_type="application/json",
+                        response_schema=list[GoogleResponse],
+                    ),
+                )
+                return response.text.strip()
+
+    async def get_response_single_chunk(
+        self, file: str, title: str, description: str, chunk: str
+    ) -> str:
+        """Get the response for a single chunk.
+
+        Args:
+            file (str): The file name.
+            title (str): The pull request title.
+            description (str): The pull request description.
+            chunk (str): The diff chunk.
+
+        Returns:
+            str: The response.
+        """
+        prompt = SINGLE_CHUNK_USER_PROMPT.format(file, title, description, chunk)
+        return await self.request(prompt)
+
+    async def get_response_full_context(
+        self, title: str, description: str, file_contents: list[str]
+    ) -> str:
+        """Get the response for full context.
+
+        Args:
+            title (str): The pull request title.
+            description (str): The pull request description.
+            file_contents (list[str]): The file contents, diffs.
+
+        Returns:
+            str: The response.
+        """
+        try:
+            prompt = FULL_CONTEXT_USER_PROMPT.format(
+                title, description, "\n".join(file_contents)
+            )
+            return await self.request(prompt)
+        except Exception as e:
+            print(f"Error during full context response: {e}")
+            print(prompt)
+            return None
+
+
+SINGLE_CHUNK_SYSTEM_PROMPT = (
+    "Your task is to review pull requests. Instructions:\n"
+    "- Provide the response in the following JSON format:  "
+    """[{{"lineNumber": int, "reviewComment": str}}] \n"""
+    "- lineNumber is about the line number of the code that in new file. \n"
+    "- lineNumber can be found at the front of each line. \n"
+    "- At the first number is old line number, the second number is new line number. \n"
+    "- If the line starts with `+`, it means the line is added. \n"
+    "- If the line starts with `-`, it means the line is deleted. \n"
+    "- Evaluate whether the code changes and additions are appropriate "
+    "and if the new code structure is suitable. \n"
+    "- Do not give positive comments or compliments. \n"
+    "- Provide comments and suggestions ONLY if there is something to improve"
+    "otherwise return an empty array. \n"
+    "- Write the comment in GitHub Markdown format. \n"
+    "- Use the given description only for the overall context "
+    "and only comment the code. \n"
+    "- Do not suggest type hint or naming convention. \n"
+    "- IMPORTANT: NEVER suggest adding comments to the code. \n"
+)
+SINGLE_CHUNK_USER_PROMPT = (
+    "Review the following code diff in the file "
+    "{} and take the pull request title and description into account "
+    "when writing the response. \n"
+    "Pull request title: {} \n"
+    "Pull request description: \n"
+    "--- \n"
+    "{} \n"
+    "--- \n"
+    "Git diff to review: \n"
+    "```diff \n"
+    "{} \n"
+    "```"
+)
+
+FULL_CONTEXT_SYSTEM_PROMPT = (
+    "You are an experienced software engineer specializing in reviewing pull "
+    "requests. Your task is to provide an overall code review summary for a PR. "
+    "Focus on assessing the following aspects:\n"
+    "1. **Code Structure & Architecture:** "
+    "Evaluate whether the code is well-organized, modular, "
+    "and adheres to clean code principles. Suggest improvements if needed.\n"
+    "2. **Refactoring Opportunities:** "
+    "Identify areas where the code can be optimized or simplified without changing "
+    "its behavior.\n"
+    "3. **Potential Future Problems:** "
+    "Highlight possible scalability, maintainability, or dependency issues that might "
+    "arise in the future based on the current implementation.\n"
+    "Be constructive and clear in your feedback. Avoid commenting on trivial issues "
+    "or syntax errors—focus on high-level feedback.\n"
+    "Precise instructions:\n"
+    "- Do not give positive comments or compliments.\n"
+    "- Provide comments and suggestions ONLY if there is something to improve, "
+    "otherwise return an empty string.\n"
+    "- Write the comment in GitHub Markdown format.\n"
+    "- Do not start with 'markdown' or '```markdown'.\n"
+    "- IMPORTANT: Give example code block or pseudo code if you can.\n"
+)
+
+FULL_CONTEXT_USER_PROMPT = (
+    "Review the following code and take the pull request title "
+    "and description into account when writing the response. \n"
+    "Pull request title: {} \n"
+    "Pull request description: \n"
+    "--- \n"
+    "{} \n"
+    "--- \n"
+    "Code to review: \n"
+    "{}"
+)

.gitea/workflows/ai-review.yml

@@ -0,0 +1,35 @@
+name: Code Review
+
+on:
+  pull_request:
+    types: [opened, synchronize]
+
+permissions:
+  contents: read
+  pull-requests: write
+
+jobs:
+  review:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install aiohttp requests py-gitea openai anthropic google-generativeai
+
+      - name: Run Code Review
+        env:
+          ACCESS_TOKEN: ${{ secrets.ACCESS_TOKEN }}
+          FULL_CONTEXT_MODEL: o3-mini
+          FULL_CONTEXT_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          SINGLE_CHUNK_MODEL: gemini-2.0-flash-exp
+          SINGLE_CHUNK_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
+          EXCLUDE: "*.yml,*.yaml"
+        run: python .gitea/scripts/code_review.py
+

.gitea/workflows/ci.yaml

@@ -0,0 +1,38 @@
+name: CI
+
+on:
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  build:
+    name: Run rust tests
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+    - uses: dtolnay/rust-toolchain@stable
+      with:
+        components: clippy
+    - uses: Swatinem/rust-cache@v2
+
+    - name: run clippy
+      run: cargo clippy -- -Dwarnings
+
+    - name: run rust tests
+      run: cargo test --verbose
+
+    - name: check compilation
+      run: cargo check
+
+  lint:
+    name: Check Rust code with rustfmt and clippy
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: dtolnay/rust-toolchain@stable
+        with:
+            components: rustfmt, clippy
+      - name: run rustfmt
+        run: cargo fmt --check
+

Cargo.toml

@@ -6,3 +6,5 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+nalgebra = "0.33.2"
+approx = "0.5"

src/context.rs

@@ -1,2 +1,20 @@
+use crate::integrator::Integrator;
+use crate::platform::Platform;
+use crate::state::{State, StateType};
+use crate::system::System;
+use nalgebra::Vector3;
 
-pub trait Context {}
+pub struct Context {
+    pub system: System,
+    pub integrator: Box<dyn Integrator>,
+    pub platform: Box<dyn Platform>,
+}
+impl Context {
+    pub fn set_positions(&mut self, _positions: Vec<Vector3<f32>>) {
+        unimplemented!()
+    }
+
+    pub fn get_state(&self, _state_type: StateType) -> State {
+        unimplemented!()
+    }
+}

src/force/harmonic_angle_force.rs

@@ -0,0 +1,179 @@
+use crate::context::Context;
+use crate::force::Force;
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct HarmonicAngleForce {}
+impl HarmonicAngleForce {
+    pub fn new() -> Self {
+        Self {}
+    }
+
+    pub fn add_angle(
+        &mut self,
+        _particle1: usize,
+        _particle2: usize,
+        _particle3: usize,
+        _angle: f64,
+        _k: f64,
+    ) {
+        unimplemented!()
+    }
+
+    pub fn set_angle_parameters(
+        &mut self,
+        _angle_id: usize,
+        _particle1: usize,
+        _particle2: usize,
+        _particle3: usize,
+        _angle: f64,
+        _k: f64,
+    ) {
+        unimplemented!()
+    }
+}
+
+impl Force for HarmonicAngleForce {
+    fn update_parameters_in_context(&mut self, _context: &mut Context) {
+        unimplemented!()
+    }
+
+    fn set_periodic_boundary_conditions(&mut self, _status: bool) {
+        unimplemented!()
+    }
+
+    fn use_periodic_boundary_conditions(&self) -> bool {
+        unimplemented!()
+    }
+}
+
+impl Default for HarmonicAngleForce {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::HarmonicAngleForce;
+    use crate::context::Context;
+    use crate::force::Force;
+    use crate::integrator::VerletIntegrator;
+    use crate::platform::CPU;
+    use crate::state::StateType;
+    use crate::system::System;
+    use approx::assert_relative_eq;
+    use nalgebra::Vector3;
+    use std::f64::consts::PI;
+
+    #[ignore]
+    #[test]
+    fn test_harmonic_angle_force() {
+        let mut system = System::new();
+        system.add_particle(1.0);
+        system.add_particle(1.0);
+        system.add_particle(1.0);
+        system.add_particle(1.0);
+        let integrator = VerletIntegrator { dt: 0.01 };
+
+        let mut forcefield = HarmonicAngleForce::new();
+        forcefield.add_angle(0, 1, 2, PI / 3.0, 1.1);
+        forcefield.add_angle(1, 2, 3, PI / 2.0, 1.2);
+
+        system.add_force(Box::new(forcefield.clone()));
+        assert_eq!(forcefield.use_periodic_boundary_conditions(), false);
+        assert_eq!(system.use_periodic_boundary_conditions(), false);
+
+        let mut context = Context {
+            system,
+            integrator: Box::new(integrator),
+            platform: Box::new(CPU),
+        };
+
+        let mut positions = vec![Vector3::zeros(); 4];
+        positions[0] = Vector3::new(0.0, 2.0, 0.0);
+        positions[1] = Vector3::new(0.0, 0.0, 0.0);
+        positions[2] = Vector3::new(1.0, 0.0, 0.0);
+        positions[3] = Vector3::new(2.0, 1.0, 0.0);
+
+        context.set_positions(positions);
+        let state = context.get_state(StateType::Either);
+        let forces = state.get_forces();
+        let torque1 = 1.1 * PI / 6.0;
+        let torque2 = 1.2 * PI / 4.0;
+        assert_relative_eq!(forces[0], Vector3::new(torque1, 0.0, 0.0));
+        assert_relative_eq!(forces[3], Vector3::new(-0.5 * torque2, 0.5 * torque2, 0.0));
+        assert_relative_eq!(forces[1], -forces[0] - forces[2]);
+        assert_relative_eq!(
+            state.get_potential_energy(),
+            0.5 * 1.1 * PI / 6.0 * PI / 6.0 + 0.5 * 1.2 * PI / 4.0 * PI / 4.0
+        );
+
+        forcefield.set_angle_parameters(0, 0, 1, 2, PI / 3.1, 1.3);
+        forcefield.set_angle_parameters(1, 1, 2, 3, PI / 2.1, 1.4);
+        forcefield.update_parameters_in_context(&mut context);
+
+        let state = context.get_state(StateType::Either);
+        let forces = state.get_forces();
+        let dtheta1 = (PI / 2.0) - (PI / 3.1);
+        let dtheta2 = (3.0 * PI / 4.0) - (PI / 2.1);
+        let torque1 = 1.3 * dtheta1;
+        let torque2 = 1.4 * dtheta2;
+        assert_relative_eq!(forces[0], Vector3::new(torque1, 0.0, 0.0));
+        assert_relative_eq!(forces[3], Vector3::new(-0.5 * torque2, 0.5 * torque2, 0.0));
+        assert_relative_eq!(forces[1], -forces[0] - forces[2]);
+        assert_relative_eq!(
+            state.get_potential_energy(),
+            0.5 * 1.3 * dtheta1 * dtheta1 + 0.5 * 1.4 * dtheta2 * dtheta2
+        );
+    }
+
+    #[ignore]
+    #[test]
+    fn test_harmonic_angle_force_with_periodic_boundary_conditions() {
+        let mut system = System {};
+        system.add_particle(1.0);
+        system.add_particle(1.0);
+        system.add_particle(1.0);
+        system.set_default_periodic_box_vectors(
+            Vector3::new(3.0, 0.0, 0.0),
+            Vector3::new(0.0, 1.5, 0.0),
+            Vector3::new(0.0, 0.0, 3.0),
+        );
+
+        let integrator = VerletIntegrator { dt: 0.01 };
+        let mut forcefield = HarmonicAngleForce {};
+        forcefield.add_angle(0, 1, 2, PI / 3.0, 1.1);
+        forcefield.set_periodic_boundary_conditions(true);
+
+        system.add_force(Box::new(forcefield));
+
+        let mut context = Context {
+            system,
+            integrator: Box::new(integrator),
+            platform: Box::new(CPU),
+        };
+
+        let mut positions = vec![Vector3::zeros(); 3];
+        positions[0] = Vector3::new(0.0, 1.0, 0.0);
+        positions[1] = Vector3::new(0.0, 0.0, 0.0);
+        positions[2] = Vector3::new(1.0, 0.0, 0.0);
+        context.set_positions(positions);
+
+        let state = context.get_state(StateType::Either);
+        let forces = state.get_forces();
+
+        let torque = 1.1 * PI / 6.0;
+        assert_relative_eq!(forces[0], Vector3::new(2.0 * torque, 0.0, 0.0));
+        assert_relative_eq!(forces[2], Vector3::new(0.0, -torque, 0.0));
+
+        let energy = state.get_potential_energy();
+        assert_relative_eq!(energy, 0.5 * 1.1 * (PI / 6.0) * (PI / 6.0));
+    }
+
+    #[ignore]
+    #[test]
+    fn test_harmonic_angle_force_parallel_computation() {
+        // TODO: Rust-OpenMM
+        unimplemented!()
+    }
+}

src/force/harmonic_bond_force.rs

@@ -0,0 +1,151 @@
+use crate::context::Context;
+use crate::force::Force;
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct HarmonicBondForce {}
+impl HarmonicBondForce {
+    pub fn new() -> Self {
+        Self {}
+    }
+
+    pub fn add_bond(&mut self, _particle1: usize, _particle2: usize, _length: f64, _k: f64) {
+        unimplemented!()
+    }
+
+    pub fn set_bond_parameters(
+        &mut self,
+        _particle1: usize,
+        _particle2: usize,
+        _length: f64,
+        _k: f64,
+    ) {
+        unimplemented!()
+    }
+}
+impl Force for HarmonicBondForce {
+    fn update_parameters_in_context(&mut self, _context: &mut Context) {
+        unimplemented!()
+    }
+
+    fn set_periodic_boundary_conditions(&mut self, _status: bool) {
+        unimplemented!()
+    }
+
+    fn use_periodic_boundary_conditions(&self) -> bool {
+        unimplemented!()
+    }
+}
+
+impl Default for HarmonicBondForce {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::HarmonicBondForce;
+    use crate::context::Context;
+    use crate::force::Force;
+    use crate::integrator::VerletIntegrator;
+    use crate::platform::CPU;
+    use crate::state::StateType;
+    use crate::system::System;
+    use approx::assert_relative_eq;
+    use nalgebra::Vector3;
+
+    #[ignore]
+    #[test]
+    fn test_harmonic_bond_force() {
+        let mut system = System {};
+        system.add_particle(1.0);
+        system.add_particle(1.0);
+        system.add_particle(1.0);
+        let integrator = VerletIntegrator { dt: 0.01 };
+
+        let mut forcefield = HarmonicBondForce {};
+        forcefield.add_bond(0, 1, 1.5, 0.8);
+        forcefield.add_bond(1, 2, 1.2, 0.7);
+
+        system.add_force(Box::new(forcefield.clone()));
+
+        let mut context = Context {
+            system,
+            integrator: Box::new(integrator),
+            platform: Box::new(CPU),
+        };
+        let mut positions = vec![Vector3::zeros(); 3];
+        positions[0] = Vector3::new(0.0, 2.0, 0.0);
+        positions[1] = Vector3::new(0.0, 0.0, 0.0);
+        positions[2] = Vector3::new(1.0, 0.0, 0.0);
+        context.set_positions(positions);
+
+        let state = context.get_state(StateType::Either);
+        let forces = state.get_forces();
+        assert_relative_eq!(forces[0], Vector3::new(0.0, -0.8 * 0.5, 0.0));
+        assert_relative_eq!(forces[2], Vector3::new(0.7 * 0.2, 0.0, 0.0));
+        assert_relative_eq!(forces[1], -forces[0] - forces[2]);
+
+        let energy = state.get_potential_energy();
+        assert_relative_eq!(energy, 0.5 * 0.8 * 0.5 * 0.5 + 0.5 * 0.7 * 0.2 * 0.2);
+
+        forcefield.set_bond_parameters(0, 0, 1.6, 0.9);
+        forcefield.set_bond_parameters(1, 1, 1.3, 0.8);
+        forcefield.update_parameters_in_context(&mut context);
+
+        let state = context.get_state(StateType::Both);
+        let forces = state.get_forces();
+        assert_relative_eq!(forces[0], Vector3::new(0.0, -0.9 * 0.4, 0.0));
+        assert_relative_eq!(forces[2], Vector3::new(0.8 * 0.3, 0.0, 0.0));
+        assert_relative_eq!(forces[1], -forces[0] - forces[2]);
+
+        let energy = state.get_potential_energy();
+        assert_relative_eq!(energy, 0.5 * 0.9 * 0.4 * 0.4 + 0.5 * 0.8 * 0.3 * 0.3);
+    }
+
+    #[ignore]
+    #[test]
+    fn test_harmonic_bond_force_with_periodic_boundary_conditions() {
+        let mut system = System {};
+        system.add_particle(1.0);
+        system.add_particle(1.0);
+        system.set_default_periodic_box_vectors(
+            Vector3::new(3.0, 0.0, 0.0),
+            Vector3::new(0.0, 3.0, 0.0),
+            Vector3::new(0.0, 0.0, 3.0),
+        );
+
+        let integrator = VerletIntegrator { dt: 0.01 };
+        let mut forcefield = HarmonicBondForce {};
+        forcefield.add_bond(0, 1, 1.2, 0.8);
+        forcefield.set_periodic_boundary_conditions(true);
+
+        system.add_force(Box::new(forcefield));
+
+        let mut context = Context {
+            system,
+            integrator: Box::new(integrator),
+            platform: Box::new(CPU),
+        };
+
+        let mut positions = vec![Vector3::zeros(); 2];
+        positions[0] = Vector3::new(0.0, 2.0, 0.0);
+        positions[1] = Vector3::new(0.0, 0.0, 0.0);
+        context.set_positions(positions);
+
+        let state = context.get_state(StateType::Both);
+        let forces = state.get_forces();
+        assert_relative_eq!(forces[0], Vector3::new(0.0, -0.8 * 0.2, 0.0));
+        assert_relative_eq!(forces[1], Vector3::new(0.0, 0.8 * 0.2, 0.0));
+
+        let energy = state.get_potential_energy();
+        assert_relative_eq!(energy, 0.5 * 0.8 * 0.2 * 0.2);
+    }
+
+    #[ignore]
+    #[test]
+    fn test_harmonic_bond_force_parallel_computation() {
+        // TODO: Rust-OpenMM
+        unimplemented!()
+    }
+}

src/force/mod.rs

@@ -0,0 +1,15 @@
+use std::fmt::Debug;
+
+use crate::context::Context;
+
+pub trait Force: Debug {
+    fn update_parameters_in_context(&mut self, context: &mut Context);
+    fn set_periodic_boundary_conditions(&mut self, status: bool);
+    fn use_periodic_boundary_conditions(&self) -> bool;
+}
+
+pub mod harmonic_angle_force;
+pub use harmonic_angle_force::HarmonicAngleForce;
+
+pub mod harmonic_bond_force;
+pub use harmonic_bond_force::HarmonicBondForce;

src/integrator/mod.rs

@@ -0,0 +1,6 @@
+pub trait Integrator {}
+
+pub struct VerletIntegrator {
+    pub dt: f64,
+}
+impl Integrator for VerletIntegrator {}

src/lib.rs

@@ -1,7 +1,8 @@
-
-
-
-pub mod system;
 pub mod context;
+pub mod force;
 pub mod platform;
 pub mod state;
+pub mod system;
+
+pub mod integrator;
+pub mod virtual_site;

src/platform.rs

@@ -1,2 +0,0 @@
-
-pub trait Platform {}

src/platform/mod.rs

@@ -0,0 +1,7 @@
+pub trait Platform {}
+
+pub struct CPU;
+pub struct OpenCL;
+
+impl Platform for CPU {}
+impl Platform for OpenCL {}

src/state.rs

@@ -1,3 +1,24 @@
+use nalgebra::Vector3;
 
+pub enum StateType {
+    Both,
+    Either,
+    Forces,
+    Energy,
+}
 
-pub trait State {}
+pub struct State {
+    pub state_type: StateType,
+    pub forces: Vec<Vector3<f64>>,
+    pub energy: f64,
+}
+
+impl State {
+    pub fn get_forces(&self) -> &Vec<Vector3<f64>> {
+        &self.forces
+    }
+
+    pub fn get_potential_energy(&self) -> f64 {
+        self.energy
+    }
+}

src/system.rs

@@ -1,3 +1,211 @@
+use crate::force::Force;
+use crate::virtual_site::VirtualSite;
+use nalgebra::Vector3;
+pub struct System {}
+impl System {
+    pub fn new() -> Self {
+        Self {}
+    }
 
+    pub fn add_particle(&mut self, _mass: f32) {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
 
-pub trait System {}
+    pub fn get_num_particles(&self) -> usize {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn get_particle_mass(&self, _index: usize) -> f32 {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn set_particle_mass(&mut self, _index: usize, _mass: f32) {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn is_virtual_site(&self, _index: usize) -> bool {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn get_virtual_site(&self, _index: usize) -> Option<Box<dyn VirtualSite>> {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn set_virtual_site(&mut self, _index: usize, _vsite: Option<Box<dyn VirtualSite>>) {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn add_constraint(&mut self, _p1: usize, _p2: usize, _dist: f32) {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn remove_constraint(&mut self, _index: usize) {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn get_num_constraints(&self) -> usize {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn get_constraint_parameters(&self, _index: usize) -> (usize, usize, f32) {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn set_constraint_parameters(&mut self, _index: usize, _p1: usize, _p2: usize, _dist: f32) {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn add_force(&mut self, _force: Box<dyn Force>) {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn remove_force(&mut self, _index: usize) {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn get_num_forces(&self) -> usize {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn get_force(&self, _index: usize) -> Box<dyn Force> {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn set_default_periodic_box_vectors(
+        &mut self,
+        _a: Vector3<f32>,
+        _b: Vector3<f32>,
+        _c: Vector3<f32>,
+    ) {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+
+    pub fn use_periodic_boundary_conditions(&self) -> bool {
+        // Todo: Rust-OpenMM
+        unimplemented!()
+    }
+}
+
+impl Default for System {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::force::{HarmonicAngleForce, HarmonicBondForce};
+    use crate::system::System;
+    use crate::virtual_site::TwoParticleAverageSite;
+
+    // Todo: Rust-OpenMM
+    #[ignore]
+    #[test]
+    fn test_system_particles() {
+        let num_particles = 10;
+        let mut system = System {};
+
+        for i in 0..num_particles {
+            system.add_particle(1.0 + 0.1 * i as f32);
+        }
+        system.set_particle_mass(5, 100.0);
+        assert_eq!(system.get_num_particles(), num_particles);
+
+        for i in 0..num_particles {
+            let expected_mass = if i == 5 { 100.0 } else { 1.0 + 0.1 * i as f32 };
+            assert_eq!(system.get_particle_mass(i), expected_mass);
+        }
+    }
+
+    // Todo: Rust-OpenMM
+    #[ignore]
+    #[test]
+    fn test_system_constraints() {
+        let num_particles = 10;
+        let mut system = System {};
+
+        for i in 0..num_particles - 1 {
+            system.add_constraint(i, i + 1, 0.2 * i as f32);
+        }
+        system.remove_constraint(5);
+        system.set_constraint_parameters(3, 0, 5, 99.0);
+        assert_eq!(system.get_num_constraints(), num_particles - 2);
+
+        for i in 0..num_particles - 2 {
+            let (p1, p2, dist) = system.get_constraint_parameters(i);
+            if i == 3 {
+                assert_eq!(p1, 0);
+                assert_eq!(p2, 5);
+                assert_eq!(dist, 99.0);
+            } else {
+                let j = if i < 5 { i } else { i + 1 };
+                assert_eq!(p1, j);
+                assert_eq!(p2, j + 1);
+                assert_eq!(dist, 0.2 * j as f32);
+            }
+        }
+    }
+
+    // Todo: Rust-OpenMM
+    #[ignore]
+    #[test]
+    fn test_system_force() {
+        let mut system = System {};
+        let bonds = HarmonicBondForce {};
+        system.add_force(Box::new(bonds));
+
+        let angles = HarmonicAngleForce {};
+        system.add_force(Box::new(angles));
+
+        assert_eq!(system.get_num_forces(), 2);
+        // Todo: Rust-OpenMM
+        // assert_eq!(system.get_force(0), Box::new(bonds));
+        // assert_eq!(system.get_force(1), Box::new(angles));
+
+        system.remove_force(0);
+        assert_eq!(system.get_num_forces(), 1);
+        // Todo: Rust-OpenMM
+        // assert_eq!(system.get_force(0), Box::new(angles));
+    }
+
+    // Todo: Rust-OpenMM
+    #[ignore]
+    #[test]
+    fn test_system_virtual_site() {
+        let mut system = System {};
+        let num_particles = 10;
+        for i in 0..num_particles {
+            system.add_particle(1.0);
+            assert!(!system.is_virtual_site(i));
+        }
+
+        let vsite = TwoParticleAverageSite::new(2, 3, 0.4, 0.6);
+        system.set_virtual_site(4, Some(Box::new(vsite)));
+        for i in 0..num_particles {
+            assert_eq!(system.is_virtual_site(i), i == 4);
+        }
+        // Todo: Rust-OpenMM
+        // assert_eq!(system.get_virtual_site(4), Some(Box::new(vsite)));
+        system.set_virtual_site(4, None);
+        for i in 0..num_particles {
+            assert!(!system.is_virtual_site(i));
+        }
+    }
+}

src/virtual_site.rs

@@ -0,0 +1,17 @@
+use std::fmt::Debug;
+
+pub trait VirtualSite: Debug {}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct TwoParticleAverageSite {
+    pub p1: usize,
+    pub p2: usize,
+    pub weight1: f32,
+    pub weight2: f32,
+}
+impl TwoParticleAverageSite {
+    pub fn new(_p1: usize, _p2: usize, _weight1: f32, _weight2: f32) -> Self {
+        unimplemented!()
+    }
+}
+impl VirtualSite for TwoParticleAverageSite {}