Program Listing for File PartioReaderWriter.cpp

Return to documentation for file (Utilities/PartioReaderWriter.cpp)

#include "PartioReaderWriter.h"
#include "extern/partio/src/lib/Partio.h"
#include "FileSystem.h"

using namespace Utilities;


bool PartioReaderWriter::readParticles(const std::string &fileName, const Vector3r &translation, const Matrix3r &rotation, const Real scale,
    std::vector<Vector3r> &positions, std::vector<Vector3r> &velocities)
{
    if (!FileSystem::fileExists(fileName))
        return false;

    Partio::ParticlesDataMutable* data = Partio::read(fileName.c_str());
    if (!data)
        return false;

    unsigned int posIndex = 0xffffffff;
    unsigned int velIndex = 0xffffffff;

    for (int i = 0; i < data->numAttributes(); i++)
    {
        Partio::ParticleAttribute attr;
        data->attributeInfo(i, attr);
        if (attr.name == "position")
            posIndex = i;
        else if (attr.name == "velocity")
            velIndex = i;
    }

    Partio::ParticleAttribute attr;

    if (posIndex != 0xffffffff)
    {
        unsigned int fSize = (unsigned int) positions.size();
        positions.resize(fSize + data->numParticles());
        data->attributeInfo(posIndex, attr);
        for (int i = 0; i < data->numParticles(); i++)
        {
            const float *pos = data->data<float>(attr, i);
            Vector3r x(pos[0], pos[1], pos[2]);
            x = rotation * (x*scale) + translation;
            positions[i + fSize] = x;
        }
    }

    if (velIndex != 0xffffffff)
    {
        unsigned int fSize = (unsigned int) velocities.size();
        velocities.resize(fSize + data->numParticles());
        data->attributeInfo(velIndex, attr);
        for (int i = 0; i < data->numParticles(); i++)
        {
            const float *vel = data->data<float>(attr, i);
            Vector3r v(vel[0], vel[1], vel[2]);
            velocities[i + fSize] = v;
        }
    }
    else
    {
        unsigned int fSize = (unsigned int) velocities.size();
        velocities.resize(fSize + data->numParticles());
        for (int i = 0; i < data->numParticles(); i++)
            velocities[i + fSize].setZero();
    }

    data->release();
    return true;
}

bool PartioReaderWriter::readParticles(const std::string &fileName, const Vector3r &translation, const Matrix3r &rotation, const Real scale,
    std::vector<Vector3r> &positions, std::vector<Vector3r> &velocities, Real &particleRadius)
{
    if (!FileSystem::fileExists(fileName))
        return false;

    Partio::ParticlesDataMutable* data = Partio::read(fileName.c_str());

    if (!data)
        return false;

    unsigned int posIndex = 0xffffffff;
    unsigned int velIndex = 0xffffffff;
    unsigned int radiusIndex = 0xffffffff;

    for (int i = 0; i < data->numAttributes(); i++)
    {
        Partio::ParticleAttribute attr;
        data->attributeInfo(i, attr);
        if (attr.name == "position")
            posIndex = i;
        else if (attr.name == "velocity")
            velIndex = i;
        else if (attr.name == "pscale")
            radiusIndex = i;
    }

    Partio::ParticleAttribute attr;

    if (posIndex != 0xffffffff)
    {
        unsigned int fSize = (unsigned int)positions.size();
        positions.resize(fSize + data->numParticles());
        data->attributeInfo(posIndex, attr);
        for (int i = 0; i < data->numParticles(); i++)
        {
            const float *pos = data->data<float>(attr, i);
            Vector3r x(pos[0], pos[1], pos[2]);
            x = rotation * (x*scale) + translation;
            positions[i + fSize] = x;
        }
    }

    if (velIndex != 0xffffffff)
    {
        unsigned int fSize = (unsigned int)velocities.size();
        velocities.resize(fSize + data->numParticles());
        data->attributeInfo(velIndex, attr);
        for (int i = 0; i < data->numParticles(); i++)
        {
            const float *vel = data->data<float>(attr, i);
            Vector3r v(vel[0], vel[1], vel[2]);
            velocities[i + fSize] = v;
        }
    }
    else
    {
        unsigned int fSize = (unsigned int)velocities.size();
        velocities.resize(fSize + data->numParticles());
        for (int i = 0; i < data->numParticles(); i++)
            velocities[i + fSize].setZero();
    }

    if (radiusIndex != 0xffffffff)
    {
        data->attributeInfo(radiusIndex, attr);
        const float *radius = data->data<float>(attr, 0);
        particleRadius = radius[0];
    }

    data->release();

    return true;
}

bool PartioReaderWriter::readParticles(const std::string &fileName, const Vector3r &translation, const Matrix3r &rotation, const Real scale,
    std::vector<Vector3r> &positions)
{
    if (!FileSystem::fileExists(fileName))
        return false;

    Partio::ParticlesDataMutable* data = Partio::read(fileName.c_str());

    if (!data)
        return false;

    unsigned int posIndex = 0xffffffff;

    for (int i = 0; i < data->numAttributes(); i++)
    {
        Partio::ParticleAttribute attr;
        data->attributeInfo(i, attr);
        if (attr.name == "position")
        {
            posIndex = i;
            break;
        }
    }

    Partio::ParticleAttribute attr;
    if (posIndex != 0xffffffff)
    {
        unsigned int fSize = (unsigned int)positions.size();
        positions.resize(fSize + data->numParticles());
        data->attributeInfo(posIndex, attr);
        for (int i = 0; i < data->numParticles(); i++)
        {
            const float *pos = data->data<float>(attr, i);
            Vector3r x(pos[0], pos[1], pos[2]);
            x = rotation * (x*scale) + translation;
            positions[i + fSize] = x;
        }
    }

    data->release();

    return true;
}


void PartioReaderWriter::writeParticles(const std::string &fileName, const unsigned int numParticles, const Vector3r *particlePositions,
    const Vector3r *particleVelocities, const Real particleRadius)
{
    if (numParticles == 0)
        return;

    Partio::ParticlesDataMutable& particleData = *Partio::create();
    Partio::ParticleAttribute posAttr = particleData.addAttribute("position", Partio::VECTOR, 3);
    Partio::ParticleAttribute velAttr;
    if (particleVelocities != NULL)
        velAttr = particleData.addAttribute("velocity", Partio::VECTOR, 3);
    Partio::ParticleAttribute scaleAttr;
    if (particleRadius != 0.0)
        scaleAttr = particleData.addAttribute("pscale", Partio::FLOAT, 1);
    Partio::ParticleAttribute idAttr = particleData.addAttribute("id", Partio::INT, 1);

    for (unsigned int i = 0; i < numParticles; i++)
    {
        Partio::ParticleIndex index = particleData.addParticle();
        float* pos = particleData.dataWrite<float>(posAttr, index);
        int* id = particleData.dataWrite<int>(idAttr, index);

        const Vector3r &x = particlePositions[i];
        pos[0] = (float)x[0];
        pos[1] = (float)x[1];
        pos[2] = (float)x[2];

        if (particleVelocities != NULL)
        {
            float* vel = particleData.dataWrite<float>(velAttr, index);
            const Vector3r &v = particleVelocities[i];
            vel[0] = (float)v[0];
            vel[1] = (float)v[1];
            vel[2] = (float)v[2];
        }

        if (particleRadius != 0.0)
        {
            float* scale = particleData.dataWrite<float>(scaleAttr, index);
            scale[0] = (float)particleRadius;
        }
        id[0] = i;
    }

    Partio::write(fileName.c_str(), particleData, true);
    particleData.release();
}